Continue accepting processes as long as the agent is active

Author: lrighi-xtm

src/main.rs

@@ -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! {

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,
             )