From cd2829c705c55533e2d627ecaf31bd1c757cbae3 Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Tue, 12 Aug 2025 12:24:17 -0400
Subject: [PATCH 01/11] Fix behavior of `--wait` when used together with
 `--agent`

---
 src/main.rs | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 6605d45..6885502 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -101,15 +101,35 @@ async fn maybe_run_until_exit(args: &argv::Args, procs: &SharedProcs) {
 
 async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
     if args.wait {
-        // Run until no processes are left, or until we receive a shutdown
-        // signal.
-        tokio::select! {
-            _ = procs.wait_idle() => {},
-            _ = procs.wait_for_shutdown() => {},
-        };
+        if args.agent {
+            // For --wait --agent: wait for exactly one run to be assigned, then disconnect
+            // First, wait for at least one process to be assigned (work arrives)
+            while procs.is_empty() {
+                // Wait for work to be assigned or shutdown signal
+                tokio::select! {
+                    _ = procs.wait_for_shutdown() => return,
+                    _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {},
+                }
+            }
 
-        // Ready to shut down now.
-        procs.set_shutdown(shutdown::State::Done);
+            // Work has been assigned! Set shutdown to Idling to prevent accepting more work
+            procs.set_shutdown(shutdown::State::Idling);
+
+            // Now wait for the single assigned process to complete and be deleted
+            tokio::select! {
+                _ = procs.wait_idle() => {},
+                _ = procs.wait_for_shutdown() => {},
+            };
+
+            // Work completed, agent should disconnect (Done state will be set by check_idling())
+        } else {
+            // Non-agent mode: just wait until idle then exit
+            tokio::select! {
+                _ = procs.wait_idle() => {},
+                _ = procs.wait_for_shutdown() => {},
+            };
+            procs.set_shutdown(shutdown::State::Done);
+        }
     }
 }
 

From dc0c5f8833c88e685cf3fae7a3c5728c0446b69b Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Tue, 12 Aug 2025 14:35:41 -0400
Subject: [PATCH 02/11] Add tests for `--wait` option

---
 tests/int/agent/test_wait_mode.py | 183 ++++++++++++++++++++++++++++++
 1 file changed, 183 insertions(+)
 create mode 100644 tests/int/agent/test_wait_mode.py

diff --git a/tests/int/agent/test_wait_mode.py b/tests/int/agent/test_wait_mode.py
new file mode 100644
index 0000000..6e5b016
--- /dev/null
+++ b/tests/int/agent/test_wait_mode.py
@@ -0,0 +1,183 @@
+import asyncio
+import pytest
+import signal
+
+from procstar import spec
+from procstar.agent.proc import Result
+from procstar.agent.exc import NoOpenConnectionInGroup
+from procstar.testing.agent import Assembly
+
+
+@pytest.mark.asyncio
+async def test_agent_selection_logic():
+    """
+    Test that the agent selection logic works correctly.
+    """
+    async with Assembly.start(args=["--wait"]) as asm:
+        available_conns = asm.server.connections._get_open_conns_in_group("default")
+
+        assert len(available_conns) == 1, (
+            f"Expected 1 available connection, got {len(available_conns)}"
+        )
+
+        conn = available_conns[0]
+        assert conn.shutdown_state.name == "active"
+
+
+@pytest.mark.asyncio
+async def test_agent_rejects_second_run():
+    """
+    Test that after accepting one run, agent becomes unavailable for more runs.
+    """
+    async with Assembly.start(args=["--wait"]) as asm:
+        conn = next(iter(asm.server.connections.values()))
+
+        # Start first process
+        proc_spec1 = spec.Proc(["/bin/sleep", "1"])
+        proc1, result1 = await asm.server.start(
+            proc_id="test-proc-1", group_id="default", spec=proc_spec1
+        )
+
+        # Give time for agent to set shutdown state to idling
+        await asyncio.sleep(0.15)
+
+        # Agent should now be in idling state, making it unavailable for new work
+        conn = asm.server.connections.get(conn.info.conn.conn_id)
+        assert conn.shutdown_state.name in ["idling"]
+
+        # Attempt to start second process should fail due to no available connections
+        proc_spec2 = spec.Proc(["/bin/echo", "second"])
+
+        with pytest.raises(NoOpenConnectionInGroup):
+            await asm.server.start(
+                proc_id="test-proc-2",
+                group_id="default",
+                spec=proc_spec2,
+                conn_timeout=0.5,  # Short timeout
+            )
+
+        # Wait for first process to complete
+        async for update in proc1.updates:
+            if isinstance(update, Result) and update.state != "running":
+                break
+
+
+@pytest.mark.asyncio
+async def test_agent_state_transitions():
+    """
+    Test the exact state transitions: active -> idling -> done.
+    """
+    async with Assembly.start(args=["--wait"]) as asm:
+        conn = next(iter(asm.server.connections.values()))
+        assert len(asm.server.connections) >= 1
+        assert conn.shutdown_state.name == "active", (
+            f"Agent should be active when idle, got: {conn.shutdown_state.name}"
+        )
+
+        await asyncio.sleep(1)
+        #  agent stays active when no work is assigned
+        assert conn.shutdown_state.name == "active", (
+            f"Agent should remain active when idle, got: {conn.shutdown_state.name}"
+        )
+
+        # Initially active
+        assert conn.shutdown_state.name == "active"
+
+        # Start a process that runs for a short time
+        proc_spec = spec.Proc(["/bin/sleep", "0.2"])
+        proc, result = await asm.server.start(
+            proc_id="test-proc", group_id="default", spec=proc_spec
+        )
+
+        # Give time for state transition to idling
+        await asyncio.sleep(0.15)
+
+        conn_id = conn.info.conn.conn_id
+        # Should now be idling (prevents new work)
+        conn = asm.server.connections.get(conn_id)
+        if conn:
+            assert conn.shutdown_state.name == "idling"
+
+        # Wait for process to complete
+        async for update in proc.updates:
+            if isinstance(update, Result) and update.state != "running":
+                break
+
+        # Give time for final state transition and cleanup
+        await asyncio.sleep(0.3)
+
+        # Connection should be done or cleaned up
+        conn = asm.server.connections.get(conn_id)
+        if conn:
+            # Agent might still be connected but should be done or idling (about to disconnect)
+            assert conn.shutdown_state.name in ["done", "idling"]
+
+
+@pytest.mark.asyncio
+async def test_multiple_agents():
+    """
+    Test multiple single-run agents can each accept one run.
+    """
+    async with Assembly.start(counts={"default": 3}, args=["--wait"]) as asm:
+        assert len(asm.server.connections) == 3
+
+        for conn in asm.server.connections.values():
+            assert conn.shutdown_state.name == "active"
+
+        # Start 3 processes - each should go to a different agent
+        procs = []
+        for i in range(3):
+            proc_spec = spec.Proc(["/bin/echo", f"task-{i}"])
+            proc, result = await asm.server.start(
+                proc_id=f"test-proc-{i}", group_id="default", spec=proc_spec
+            )
+            procs.append(proc)
+            # Give time for agent state changes to propagate
+            await asyncio.sleep(0.2)
+
+        # Give time for state transitions
+        await asyncio.sleep(0.8)
+
+        # All agents should now be idling (no longer accepting work)
+        active_count = sum(
+            1 for conn in asm.server.connections.values() if conn.shutdown_state.name == "active"
+        )
+        assert active_count == 0, "All agents should have transitioned from active state"
+
+        # Wait for all processes to complete
+        for proc in procs:
+            async for update in proc.updates:
+                if isinstance(update, Result) and update.state != "running":
+                    break
+
+        # Give time for cleanup
+        await asyncio.sleep(0.1)
+
+        with pytest.raises(NoOpenConnectionInGroup):
+            proc_spec = spec.Proc(["/bin/echo", "should-fail"])
+            await asm.server.start(
+                proc_id="test-proc-fail",
+                group_id="default",
+                spec=proc_spec,
+                conn_timeout=0.1,
+            )
+
+
+@pytest.mark.asyncio
+async def test_agent_shutdown_before_receiving_work():
+    """
+    Test that agent responds to shutdown signals even before receiving work.
+    """
+    async with Assembly.start(args=["--wait"]) as asm:
+        conn = next(iter(asm.server.connections.values()))
+        conn_id = conn.info.conn.conn_id
+
+        # Send shutdown signal to the procstar process
+        procstar_proc = asm.conn_procs[conn_id]
+        procstar_proc.send_signal(signal.SIGUSR1)
+
+        # Wait for graceful shutdown
+        await asyncio.sleep(0.5)
+
+        # Connection should be cleaned up
+        assert conn_id not in asm.server.connections

From 86ac4d5a0c24840a1f0b7b6e13606d2b7723c37d Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Wed, 13 Aug 2025 06:06:36 -0400
Subject: [PATCH 03/11] Minor refactor to reduce nesting

---
 src/main.rs | 53 +++++++++++++++++++++++++++--------------------------
 1 file changed, 27 insertions(+), 26 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 6885502..613ef1e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -100,36 +100,37 @@ async fn maybe_run_until_exit(args: &argv::Args, procs: &SharedProcs) {
 }
 
 async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
-    if args.wait {
-        if args.agent {
-            // For --wait --agent: wait for exactly one run to be assigned, then disconnect
-            // First, wait for at least one process to be assigned (work arrives)
-            while procs.is_empty() {
-                // Wait for work to be assigned or shutdown signal
-                tokio::select! {
-                    _ = procs.wait_for_shutdown() => return,
-                    _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {},
-                }
+    if !args.wait {
+        return;
+    }
+    if args.agent {
+        // For --wait --agent: wait for exactly one run to be assigned, then disconnect
+        // First, wait for at least one process to be assigned (work arrives)
+        while procs.is_empty() {
+            // Wait for work to be assigned or shutdown signal
+            tokio::select! {
+                _ = procs.wait_for_shutdown() => return,
+                _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {},
             }
+        }
 
-            // Work has been assigned! Set shutdown to Idling to prevent accepting more work
-            procs.set_shutdown(shutdown::State::Idling);
+        // Work has been assigned! Set shutdown to Idling to prevent accepting more work
+        procs.set_shutdown(shutdown::State::Idling);
 
-            // Now wait for the single assigned process to complete and be deleted
-            tokio::select! {
-                _ = procs.wait_idle() => {},
-                _ = procs.wait_for_shutdown() => {},
-            };
+        // Now wait for the single assigned process to complete and be deleted
+        tokio::select! {
+            _ = procs.wait_idle() => {},
+            _ = procs.wait_for_shutdown() => {},
+        };
 
-            // Work completed, agent should disconnect (Done state will be set by check_idling())
-        } else {
-            // Non-agent mode: just wait until idle then exit
-            tokio::select! {
-                _ = procs.wait_idle() => {},
-                _ = procs.wait_for_shutdown() => {},
-            };
-            procs.set_shutdown(shutdown::State::Done);
-        }
+        // Work completed, agent should disconnect (Done state will be set by check_idling())
+    } else {
+        // Non-agent mode: just wait until idle then exit
+        tokio::select! {
+            _ = procs.wait_idle() => {},
+            _ = procs.wait_for_shutdown() => {},
+        };
+        procs.set_shutdown(shutdown::State::Done);
     }
 }
 

From 7a8516ea362f99acfd0253927891841b2935f596 Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Wed, 13 Aug 2025 06:25:03 -0400
Subject: [PATCH 04/11] Replace polling with notification mechanism

---
 src/main.rs | 32 +++++++++++++++++++++-----------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 613ef1e..89005da 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,7 +6,7 @@ use log::*;
 // use procstar::fd::parse_fd;
 use procstar::agent;
 use procstar::http;
-use procstar::procs::{restrict_exe, start_procs, SharedProcs};
+use procstar::procs::{restrict_exe, start_procs, SharedProcs, Notification};
 use procstar::proto;
 use procstar::res;
 use procstar::shutdown;
@@ -103,27 +103,37 @@ async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
     if !args.wait {
         return;
     }
+
     if args.agent {
         // For --wait --agent: wait for exactly one run to be assigned, then disconnect
-        // First, wait for at least one process to be assigned (work arrives)
-        while procs.is_empty() {
-            // Wait for work to be assigned or shutdown signal
-            tokio::select! {
-                _ = procs.wait_for_shutdown() => return,
-                _ = tokio::time::sleep(tokio::time::Duration::from_millis(100)) => {},
+        let mut sub = procs.subscribe();
+        
+        // If already has work, skip waiting
+        if procs.is_empty() {
+            // Wait for first process assignment
+            loop {
+                tokio::select! {
+                    _ = procs.wait_for_shutdown() => return,
+                    notification = sub.recv() => {
+                        match notification {
+                            Some(Notification::Start(_)) => break,
+                            Some(_) => continue,
+                            None => return,
+                        }
+                    }
+                }
             }
         }
 
-        // Work has been assigned! Set shutdown to Idling to prevent accepting more work
+        // Process has been assigned! Set shutdown to Idling to prevent accepting other processes
         procs.set_shutdown(shutdown::State::Idling);
 
-        // Now wait for the single assigned process to complete and be deleted
+        // Now wait for the single assigned process to complete
         tokio::select! {
             _ = procs.wait_idle() => {},
             _ = procs.wait_for_shutdown() => {},
         };
-
-        // Work completed, agent should disconnect (Done state will be set by check_idling())
+        // Process completed and deleted, agent should disconnect
     } else {
         // Non-agent mode: just wait until idle then exit
         tokio::select! {

From 3014e3895400b5cda08a86a64f5f41091c2db789 Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Wed, 13 Aug 2025 10:21:46 -0400
Subject: [PATCH 05/11] Continue accepting processes as long as the agent is
 active

---
 src/main.rs                       |  17 ++---
 tests/int/agent/test_wait_mode.py | 120 ++++--------------------------
 2 files changed, 24 insertions(+), 113 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 89005da..1f88c4b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,7 +6,7 @@ use log::*;
 // use procstar::fd::parse_fd;
 use procstar::agent;
 use procstar::http;
-use procstar::procs::{restrict_exe, start_procs, SharedProcs, Notification};
+use procstar::procs::{restrict_exe, start_procs, Notification, SharedProcs};
 use procstar::proto;
 use procstar::res;
 use procstar::shutdown;
@@ -105,10 +105,11 @@ async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
     }
 
     if args.agent {
-        // For --wait --agent: wait for exactly one run to be assigned, then disconnect
+        // For --wait --agent: wait for at least one process to be assigned,
+        // then wait for all processes to be deleted before disconnecting
         let mut sub = procs.subscribe();
-        
-        // If already has work, skip waiting
+
+        // If already has work, skip waiting for first assignment
         if procs.is_empty() {
             // Wait for first process assignment
             loop {
@@ -125,15 +126,13 @@ async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
             }
         }
 
-        // Process has been assigned! Set shutdown to Idling to prevent accepting other processes
-        procs.set_shutdown(shutdown::State::Idling);
-
-        // Now wait for the single assigned process to complete
+        // At least one process has been assigned! Continue accepting processes
+        // and only shutdown when all assigned processes are deleted
         tokio::select! {
             _ = procs.wait_idle() => {},
             _ = procs.wait_for_shutdown() => {},
         };
-        // Process completed and deleted, agent should disconnect
+        procs.set_shutdown(shutdown::State::Done);
     } else {
         // Non-agent mode: just wait until idle then exit
         tokio::select! {
diff --git a/tests/int/agent/test_wait_mode.py b/tests/int/agent/test_wait_mode.py
index 6e5b016..20dbdd9 100644
--- a/tests/int/agent/test_wait_mode.py
+++ b/tests/int/agent/test_wait_mode.py
@@ -24,49 +24,8 @@ async def test_agent_selection_logic():
         assert conn.shutdown_state.name == "active"
 
 
-@pytest.mark.asyncio
-async def test_agent_rejects_second_run():
-    """
-    Test that after accepting one run, agent becomes unavailable for more runs.
-    """
-    async with Assembly.start(args=["--wait"]) as asm:
-        conn = next(iter(asm.server.connections.values()))
-
-        # Start first process
-        proc_spec1 = spec.Proc(["/bin/sleep", "1"])
-        proc1, result1 = await asm.server.start(
-            proc_id="test-proc-1", group_id="default", spec=proc_spec1
-        )
-
-        # Give time for agent to set shutdown state to idling
-        await asyncio.sleep(0.15)
-
-        # Agent should now be in idling state, making it unavailable for new work
-        conn = asm.server.connections.get(conn.info.conn.conn_id)
-        assert conn.shutdown_state.name in ["idling"]
-
-        # Attempt to start second process should fail due to no available connections
-        proc_spec2 = spec.Proc(["/bin/echo", "second"])
-
-        with pytest.raises(NoOpenConnectionInGroup):
-            await asm.server.start(
-                proc_id="test-proc-2",
-                group_id="default",
-                spec=proc_spec2,
-                conn_timeout=0.5,  # Short timeout
-            )
-
-        # Wait for first process to complete
-        async for update in proc1.updates:
-            if isinstance(update, Result) and update.state != "running":
-                break
-
-
 @pytest.mark.asyncio
 async def test_agent_state_transitions():
-    """
-    Test the exact state transitions: active -> idling -> done.
-    """
     async with Assembly.start(args=["--wait"]) as asm:
         conn = next(iter(asm.server.connections.values()))
         assert len(asm.server.connections) >= 1
@@ -75,91 +34,44 @@ async def test_agent_state_transitions():
         )
 
         await asyncio.sleep(1)
-        #  agent stays active when no work is assigned
+        #  agent stays active until no work is assigned
         assert conn.shutdown_state.name == "active", (
-            f"Agent should remain active when idle, got: {conn.shutdown_state.name}"
+            f"Agent should remain active until some work is assigned, got: {conn.shutdown_state.name}"
         )
 
-        # Initially active
-        assert conn.shutdown_state.name == "active"
-
         # Start a process that runs for a short time
         proc_spec = spec.Proc(["/bin/sleep", "0.2"])
         proc, result = await asm.server.start(
             proc_id="test-proc", group_id="default", spec=proc_spec
         )
 
-        # Give time for state transition to idling
-        await asyncio.sleep(0.15)
-
-        conn_id = conn.info.conn.conn_id
-        # Should now be idling (prevents new work)
-        conn = asm.server.connections.get(conn_id)
-        if conn:
-            assert conn.shutdown_state.name == "idling"
-
-        # Wait for process to complete
         async for update in proc.updates:
             if isinstance(update, Result) and update.state != "running":
                 break
 
-        # Give time for final state transition and cleanup
-        await asyncio.sleep(0.3)
-
-        # Connection should be done or cleaned up
+        conn_id = conn.info.conn.conn_id
         conn = asm.server.connections.get(conn_id)
-        if conn:
-            # Agent might still be connected but should be done or idling (about to disconnect)
-            assert conn.shutdown_state.name in ["done", "idling"]
-
-
-@pytest.mark.asyncio
-async def test_multiple_agents():
-    """
-    Test multiple single-run agents can each accept one run.
-    """
-    async with Assembly.start(counts={"default": 3}, args=["--wait"]) as asm:
-        assert len(asm.server.connections) == 3
-
-        for conn in asm.server.connections.values():
-            assert conn.shutdown_state.name == "active"
-
-        # Start 3 processes - each should go to a different agent
-        procs = []
-        for i in range(3):
-            proc_spec = spec.Proc(["/bin/echo", f"task-{i}"])
-            proc, result = await asm.server.start(
-                proc_id=f"test-proc-{i}", group_id="default", spec=proc_spec
-            )
-            procs.append(proc)
-            # Give time for agent state changes to propagate
-            await asyncio.sleep(0.2)
+        assert conn.shutdown_state.name == "active", (
+            "Agent should stay active until all processes are deleted"
+        )
 
-        # Give time for state transitions
-        await asyncio.sleep(0.8)
+        # Delete the completed process to trigger the shutdown sequence
+        await proc.delete()
 
-        # All agents should now be idling (no longer accepting work)
-        active_count = sum(
-            1 for conn in asm.server.connections.values() if conn.shutdown_state.name == "active"
-        )
-        assert active_count == 0, "All agents should have transitioned from active state"
+        await asyncio.sleep(1)
 
-        # Wait for all processes to complete
-        for proc in procs:
-            async for update in proc.updates:
-                if isinstance(update, Result) and update.state != "running":
-                    break
+        active_conns = asm.server.connections._get_open_conns_in_group("default")
+        assert len(active_conns) == 0, f"Expected no active connections, got {len(active_conns)}"
 
-        # Give time for cleanup
-        await asyncio.sleep(0.1)
+        # Attempt to start second process should fail due to no available connections
+        proc_spec2 = spec.Proc(["/bin/echo", "second"])
 
         with pytest.raises(NoOpenConnectionInGroup):
-            proc_spec = spec.Proc(["/bin/echo", "should-fail"])
             await asm.server.start(
-                proc_id="test-proc-fail",
+                proc_id="test-proc-2",
                 group_id="default",
-                spec=proc_spec,
-                conn_timeout=0.1,
+                spec=proc_spec2,
+                conn_timeout=0.5,
             )
 
 

From 23d520cb7c1a564fe30284ef99b33408f4259b3f Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Wed, 13 Aug 2025 10:49:08 -0400
Subject: [PATCH 06/11] Add a configurable timeout for the `--wait` mode

---
 src/argv.rs                       |  3 +++
 src/main.rs                       | 13 ++++++++++++-
 tests/int/agent/test_wait_mode.py | 20 ++++++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/src/argv.rs b/src/argv.rs
index 6fa3b58..7b9aee8 100644
--- a/src/argv.rs
+++ b/src/argv.rs
@@ -45,6 +45,9 @@ pub struct Args {
     /// When idle (all processes deleted), exit
     #[arg(short = 'w', long)]
     pub wait: bool,
+    /// Timeout in seconds for wait mode when no work is assigned
+    #[arg(long, value_name = "SECS", default_value_t = 600)]
+    pub wait_timeout: u64,
 
     /// Run an HTTP service
     #[arg(short, long)]
diff --git a/src/main.rs b/src/main.rs
index 1f88c4b..c62f981 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,6 +15,8 @@ use procstar::sig::{SIGINT, SIGQUIT, SIGTERM, SIGUSR1};
 use procstar::spec;
 use procstar::systemd::api::{maybe_connect, SharedSystemdClient};
 use std::rc::Rc;
+use std::time::Duration;
+use tokio::time::sleep;
 
 //------------------------------------------------------------------------------
 
@@ -111,10 +113,19 @@ async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
 
         // If already has work, skip waiting for first assignment
         if procs.is_empty() {
-            // Wait for first process assignment
+            // Wait for first process assignment with configurable timeout
+            let timeout_duration = Duration::from_secs(args.wait_timeout);
+            let timeout_sleep = sleep(timeout_duration);
+            tokio::pin!(timeout_sleep);
+
             loop {
                 tokio::select! {
                     _ = procs.wait_for_shutdown() => return,
+                    _ = &mut timeout_sleep => {
+                        warn!("agent timeout: no work assigned after {} seconds, shutting down", args.wait_timeout);
+                        procs.set_shutdown(shutdown::State::Done);
+                        return;
+                    },
                     notification = sub.recv() => {
                         match notification {
                             Some(Notification::Start(_)) => break,
diff --git a/tests/int/agent/test_wait_mode.py b/tests/int/agent/test_wait_mode.py
index 20dbdd9..f3693e8 100644
--- a/tests/int/agent/test_wait_mode.py
+++ b/tests/int/agent/test_wait_mode.py
@@ -93,3 +93,23 @@ async def test_agent_shutdown_before_receiving_work():
 
         # Connection should be cleaned up
         assert conn_id not in asm.server.connections
+
+
+@pytest.mark.asyncio
+async def test_agent_timeout_no_work():
+    """
+    Test that agent times out and shuts down when no work is assigned within the timeout period.
+    """
+    async with Assembly.start(args=["--wait", "--wait-timeout", "1"]) as asm:
+        # Verify agent starts with active connection
+        active_conns = asm.server.connections._get_open_conns_in_group("default")
+        assert len(active_conns) == 1, f"Expected 1 active connection, got {len(active_conns)}"
+
+        # Wait for timeout
+        await asyncio.sleep(1)
+
+        # Agent should have shut down due to timeout - no more active connections
+        active_conns = asm.server.connections._get_open_conns_in_group("default")
+        assert len(active_conns) == 0, (
+            f"Expected no active connections after timeout, got {len(active_conns)}"
+        )

From 4de8fa7605d715fd3a4b986035559b5ea293ac3c Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Wed, 13 Aug 2025 13:46:10 -0400
Subject: [PATCH 07/11] Refactor: split into different functions

---
 src/main.rs | 94 +++++++++++++++++++++++++++++------------------------
 1 file changed, 51 insertions(+), 43 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index c62f981..6bec1a2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -101,56 +101,64 @@ async fn maybe_run_until_exit(args: &argv::Args, procs: &SharedProcs) {
     }
 }
 
-async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
-    if !args.wait {
-        return;
-    }
+async fn run_agent_until_idle(args: &argv::Args, procs: &SharedProcs) {
+    // For --wait --agent: wait for at least one process to be assigned,
+    // then wait for all processes to be deleted before disconnecting
+    let mut sub = procs.subscribe();
 
-    if args.agent {
-        // For --wait --agent: wait for at least one process to be assigned,
-        // then wait for all processes to be deleted before disconnecting
-        let mut sub = procs.subscribe();
-
-        // If already has work, skip waiting for first assignment
-        if procs.is_empty() {
-            // Wait for first process assignment with configurable timeout
-            let timeout_duration = Duration::from_secs(args.wait_timeout);
-            let timeout_sleep = sleep(timeout_duration);
-            tokio::pin!(timeout_sleep);
-
-            loop {
-                tokio::select! {
-                    _ = procs.wait_for_shutdown() => return,
-                    _ = &mut timeout_sleep => {
-                        warn!("agent timeout: no work assigned after {} seconds, shutting down", args.wait_timeout);
-                        procs.set_shutdown(shutdown::State::Done);
-                        return;
-                    },
-                    notification = sub.recv() => {
-                        match notification {
-                            Some(Notification::Start(_)) => break,
-                            Some(_) => continue,
-                            None => return,
-                        }
+    // If already has work, skip waiting for first assignment
+    if procs.is_empty() {
+        // Wait for first process assignment with configurable timeout
+        let timeout_duration = Duration::from_secs(args.wait_timeout);
+        let timeout_sleep = sleep(timeout_duration);
+        tokio::pin!(timeout_sleep);
+
+        loop {
+            tokio::select! {
+                _ = procs.wait_for_shutdown() => return,
+                _ = &mut timeout_sleep => {
+                    warn!("agent timeout: no work assigned after {} seconds, shutting down", args.wait_timeout);
+                    procs.set_shutdown(shutdown::State::Done);
+                    return;
+                },
+                notification = sub.recv() => {
+                    match notification {
+                        Some(Notification::Start(_)) => break,
+                        Some(_) => continue,
+                        None => return,
                     }
                 }
             }
         }
+    }
 
-        // At least one process has been assigned! Continue accepting processes
-        // and only shutdown when all assigned processes are deleted
-        tokio::select! {
-            _ = procs.wait_idle() => {},
-            _ = procs.wait_for_shutdown() => {},
-        };
-        procs.set_shutdown(shutdown::State::Done);
+    // At least one process has been assigned! Continue accepting processes
+    // and only shutdown when all assigned processes are deleted
+    tokio::select! {
+        _ = procs.wait_idle() => {},
+        _ = procs.wait_for_shutdown() => {},
+    };
+    procs.set_shutdown(shutdown::State::Done);
+}
+
+async fn run_standalone_until_idle(procs: &SharedProcs) {
+    // Non-agent mode: just wait until idle then exit
+    tokio::select! {
+        _ = procs.wait_idle() => {},
+        _ = procs.wait_for_shutdown() => {},
+    };
+    procs.set_shutdown(shutdown::State::Done);
+}
+
+async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
+    if !args.wait {
+        return;
+    }
+
+    if args.agent {
+        run_agent_until_idle(args, procs).await;
     } else {
-        // Non-agent mode: just wait until idle then exit
-        tokio::select! {
-            _ = procs.wait_idle() => {},
-            _ = procs.wait_for_shutdown() => {},
-        };
-        procs.set_shutdown(shutdown::State::Done);
+        run_standalone_until_idle(procs).await;
     }
 }
 

From b60df5663188e7fac072962c92689b0fcb0f2dce Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Wed, 13 Aug 2025 14:44:51 -0400
Subject: [PATCH 08/11] General refactoring to improve readability

---
 src/main.rs | 85 +++++++++++++++++++++++++++++------------------------
 1 file changed, 46 insertions(+), 39 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 6bec1a2..3a68a60 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,7 +6,7 @@ use log::*;
 // use procstar::fd::parse_fd;
 use procstar::agent;
 use procstar::http;
-use procstar::procs::{restrict_exe, start_procs, Notification, SharedProcs};
+use procstar::procs::{restrict_exe, start_procs, Notification, NotificationSub, SharedProcs};
 use procstar::proto;
 use procstar::res;
 use procstar::shutdown;
@@ -16,7 +16,14 @@ use procstar::spec;
 use procstar::systemd::api::{maybe_connect, SharedSystemdClient};
 use std::rc::Rc;
 use std::time::Duration;
-use tokio::time::sleep;
+
+//------------------------------------------------------------------------------
+
+#[derive(Debug)]
+enum WaitError {
+    Timeout,
+    Shutdown,
+}
 
 //------------------------------------------------------------------------------
 
@@ -101,39 +108,34 @@ async fn maybe_run_until_exit(args: &argv::Args, procs: &SharedProcs) {
     }
 }
 
-async fn run_agent_until_idle(args: &argv::Args, procs: &SharedProcs) {
-    // For --wait --agent: wait for at least one process to be assigned,
-    // then wait for all processes to be deleted before disconnecting
-    let mut sub = procs.subscribe();
-
-    // If already has work, skip waiting for first assignment
-    if procs.is_empty() {
-        // Wait for first process assignment with configurable timeout
-        let timeout_duration = Duration::from_secs(args.wait_timeout);
-        let timeout_sleep = sleep(timeout_duration);
-        tokio::pin!(timeout_sleep);
-
-        loop {
-            tokio::select! {
-                _ = procs.wait_for_shutdown() => return,
-                _ = &mut timeout_sleep => {
-                    warn!("agent timeout: no work assigned after {} seconds, shutting down", args.wait_timeout);
-                    procs.set_shutdown(shutdown::State::Done);
-                    return;
-                },
-                notification = sub.recv() => {
-                    match notification {
-                        Some(Notification::Start(_)) => break,
-                        Some(_) => continue,
-                        None => return,
-                    }
+async fn wait_for_first_assignment(
+    procs: &SharedProcs,
+    mut sub: NotificationSub,
+    timeout_secs: u64,
+) -> Result<(), WaitError> {
+    let timeout_duration = Duration::from_secs(timeout_secs);
+    tokio::select! {
+        _ = tokio::time::sleep(timeout_duration) => {
+            warn!(
+                "agent timeout: no work assigned after {} seconds, shutting down",
+                timeout_secs
+            );
+            procs.set_shutdown(shutdown::State::Done);
+            Err(WaitError::Timeout)
+        }
+        _ = procs.wait_for_shutdown() => Err(WaitError::Shutdown),
+        result = async {
+            while let Some(notification) = sub.recv().await {
+                if let Notification::Start(_) = notification {
+                    return Ok(());
                 }
             }
-        }
+            Err(WaitError::Shutdown) // channel closed
+        } => result,
     }
+}
 
-    // At least one process has been assigned! Continue accepting processes
-    // and only shutdown when all assigned processes are deleted
+async fn wait_until_idle_then_shutdown(procs: &SharedProcs) {
     tokio::select! {
         _ = procs.wait_idle() => {},
         _ = procs.wait_for_shutdown() => {},
@@ -141,13 +143,18 @@ async fn run_agent_until_idle(args: &argv::Args, procs: &SharedProcs) {
     procs.set_shutdown(shutdown::State::Done);
 }
 
-async fn run_standalone_until_idle(procs: &SharedProcs) {
-    // Non-agent mode: just wait until idle then exit
-    tokio::select! {
-        _ = procs.wait_idle() => {},
-        _ = procs.wait_for_shutdown() => {},
-    };
-    procs.set_shutdown(shutdown::State::Done);
+async fn run_agent_until_idle(args: &argv::Args, procs: &SharedProcs) {
+    if procs.is_empty() {
+        let sub = procs.subscribe();
+        if wait_for_first_assignment(procs, sub, args.wait_timeout)
+            .await
+            .is_err()
+        {
+            return; // early exit on timeout or shutdown
+        }
+    }
+    // Wait until idle or shutdown
+    wait_until_idle_then_shutdown(procs).await;
 }
 
 async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
@@ -158,7 +165,7 @@ async fn maybe_run_until_idle(args: &argv::Args, procs: &SharedProcs) {
     if args.agent {
         run_agent_until_idle(args, procs).await;
     } else {
-        run_standalone_until_idle(procs).await;
+        wait_until_idle_then_shutdown(procs).await;
     }
 }
 

From 3ed36ff4a56e997c99a737aa655922a2756acdfd Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Thu, 14 Aug 2025 05:59:10 -0400
Subject: [PATCH 09/11] Refactor: use `tokio::time::timeout`

---
 src/main.rs | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 3a68a60..aafce5a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -114,8 +114,25 @@ async fn wait_for_first_assignment(
     timeout_secs: u64,
 ) -> Result<(), WaitError> {
     let timeout_duration = Duration::from_secs(timeout_secs);
-    tokio::select! {
-        _ = tokio::time::sleep(timeout_duration) => {
+
+    let work_future = async {
+        loop {
+            tokio::select! {
+                _ = procs.wait_for_shutdown() => return Err(WaitError::Shutdown),
+                notification = sub.recv() => {
+                    match notification {
+                        Some(Notification::Start(_)) => return Ok(()),
+                        Some(_) => continue,  // ignore other notification types
+                        None => return Err(WaitError::Shutdown), // channel closed
+                    }
+                }
+            }
+        }
+    };
+
+    match tokio::time::timeout(timeout_duration, work_future).await {
+        Ok(result) => result,
+        Err(_) => {
             warn!(
                 "agent timeout: no work assigned after {} seconds, shutting down",
                 timeout_secs
@@ -123,15 +140,6 @@ async fn wait_for_first_assignment(
             procs.set_shutdown(shutdown::State::Done);
             Err(WaitError::Timeout)
         }
-        _ = procs.wait_for_shutdown() => Err(WaitError::Shutdown),
-        result = async {
-            while let Some(notification) = sub.recv().await {
-                if let Notification::Start(_) = notification {
-                    return Ok(());
-                }
-            }
-            Err(WaitError::Shutdown) // channel closed
-        } => result,
     }
 }
 

From eb74eca21ee009dec5a37d1e7dfcdaa06fe85c38 Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Thu, 14 Aug 2025 06:22:31 -0400
Subject: [PATCH 10/11] Timeout only the wait for work, not the wait for signal

---
 src/main.rs | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index aafce5a..5e07b42 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -117,28 +117,28 @@ async fn wait_for_first_assignment(
 
     let work_future = async {
         loop {
-            tokio::select! {
-                _ = procs.wait_for_shutdown() => return Err(WaitError::Shutdown),
-                notification = sub.recv() => {
-                    match notification {
-                        Some(Notification::Start(_)) => return Ok(()),
-                        Some(_) => continue,  // ignore other notification types
-                        None => return Err(WaitError::Shutdown), // channel closed
-                    }
-                }
+            match sub.recv().await {
+                Some(Notification::Start(_)) => return Ok(()),
+                Some(_) => continue,
+                None => return Err(WaitError::Shutdown),
             }
         }
     };
 
-    match tokio::time::timeout(timeout_duration, work_future).await {
-        Ok(result) => result,
-        Err(_) => {
-            warn!(
-                "agent timeout: no work assigned after {} seconds, shutting down",
-                timeout_secs
-            );
-            procs.set_shutdown(shutdown::State::Done);
-            Err(WaitError::Timeout)
+    tokio::select! {
+        _ = procs.wait_for_shutdown() => Err(WaitError::Shutdown),
+        result = tokio::time::timeout(timeout_duration, work_future) => {
+            match result {
+                Ok(r) => r,
+                Err(_) => {
+                    warn!(
+                        "agent timeout: no work assigned after {} seconds, shutting down",
+                        timeout_secs
+                    );
+                    procs.set_shutdown(shutdown::State::Done);
+                    Err(WaitError::Timeout)
+                }
+            }
         }
     }
 }

From bcf50611cb277c56d57b3ad5a6ba6adce37aca0e Mon Sep 17 00:00:00 2001
From: Ludovico Righi <ludovico.righi@xantium.com>
Date: Thu, 14 Aug 2025 06:54:26 -0400
Subject: [PATCH 11/11] Simplify further

---
 src/main.rs | 65 ++++++++++++++++-------------------------------------
 1 file changed, 19 insertions(+), 46 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 5e07b42..e88638e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,14 +19,6 @@ use std::time::Duration;
 
 //------------------------------------------------------------------------------
 
-#[derive(Debug)]
-enum WaitError {
-    Timeout,
-    Shutdown,
-}
-
-//------------------------------------------------------------------------------
-
 async fn maybe_run_http(
     args: &argv::Args,
     procs: &SharedProcs,
@@ -108,37 +100,10 @@ async fn maybe_run_until_exit(args: &argv::Args, procs: &SharedProcs) {
     }
 }
 
-async fn wait_for_first_assignment(
-    procs: &SharedProcs,
-    mut sub: NotificationSub,
-    timeout_secs: u64,
-) -> Result<(), WaitError> {
-    let timeout_duration = Duration::from_secs(timeout_secs);
-
-    let work_future = async {
-        loop {
-            match sub.recv().await {
-                Some(Notification::Start(_)) => return Ok(()),
-                Some(_) => continue,
-                None => return Err(WaitError::Shutdown),
-            }
-        }
-    };
-
-    tokio::select! {
-        _ = procs.wait_for_shutdown() => Err(WaitError::Shutdown),
-        result = tokio::time::timeout(timeout_duration, work_future) => {
-            match result {
-                Ok(r) => r,
-                Err(_) => {
-                    warn!(
-                        "agent timeout: no work assigned after {} seconds, shutting down",
-                        timeout_secs
-                    );
-                    procs.set_shutdown(shutdown::State::Done);
-                    Err(WaitError::Timeout)
-                }
-            }
+async fn wait_for_first_assignment(mut sub: NotificationSub) {
+    while let Some(notification) = sub.recv().await {
+        if let Notification::Start(_) = notification {
+            return;
         }
     }
 }
@@ -153,15 +118,23 @@ async fn wait_until_idle_then_shutdown(procs: &SharedProcs) {
 
 async fn run_agent_until_idle(args: &argv::Args, procs: &SharedProcs) {
     if procs.is_empty() {
-        let sub = procs.subscribe();
-        if wait_for_first_assignment(procs, sub, args.wait_timeout)
-            .await
-            .is_err()
-        {
-            return; // early exit on timeout or shutdown
+        tokio::select! {
+            _ = procs.wait_for_shutdown() => return,
+            result = tokio::time::timeout(
+                Duration::from_secs(args.wait_timeout),
+                wait_for_first_assignment(procs.subscribe()),
+            ) => {
+                if result.is_err() {
+                    warn!(
+                        "agent timeout: no work assigned after {} seconds, shutting down",
+                        args.wait_timeout
+                    );
+                    procs.set_shutdown(shutdown::State::Done);
+                    return;
+                }
+            }
         }
     }
-    // Wait until idle or shutdown
     wait_until_idle_then_shutdown(procs).await;
 }