Fix two test_wait flakes: a stop that waited on background commands, and a test lock that held the service loop

docs/RUST_PROTOTYPE.md

@@ -251,7 +251,8 @@ Parked turns survive a daemon restart: they are re-registered at
 startup, peer handles resolve from durable state (a peer interrupted by the
 restart reports `interrupted`), and commands that were still running resolve
 to `process_lost`: their supervision ended with the daemon, whether or not the
-OS process did. Interrupting a parked
+OS process did. A daemon that stops cleanly does not wait for them: it kills
+each one's process group and records `process_lost` itself. Interrupting a parked
 turn ends it as `interrupted`: the wait and any planned calls behind it get a
 `cancelled` tool result in the same transaction, so the conversation stays
 valid and the bot is not left uncertain. A bot whose turn is parked reports

src/server/mod.rs

@@ -1596,8 +1596,10 @@ pub async fn run(config: Configuration) -> Result<()> {
     }
     service.replays.abort_all();
     while service.replays.join_next().await.is_some() {}
-    // Terminals end with the daemon, and record how on their handles.
+    // Terminals and background commands end with the daemon, and record
+    // how on their handles, so no result task keeps the store open below.
     service.registry.end_terminals();
+    service.registry.end_background();
     // Every completion is committed. Allow publication to drain before
     // releasing waiters. Background result tasks may still own Store clones,
     // keeping the stream open even after the service is dropped.

src/tools.rs

@@ -18,7 +18,7 @@ use std::{
 };
 use tokio::{
     io::{AsyncRead, AsyncReadExt},
-    sync::Semaphore,
+    sync::{Semaphore, watch},
 };
 
 mod fetch;
@@ -412,6 +412,9 @@ pub struct Registry {
     /// Running terminals, by proc handle, each its bot's own.
     #[cfg(unix)]
     terminals: terminal::Terminals,
+    /// Set once as the daemon stops: background commands still running
+    /// end then rather than outlive it.
+    stopping: Arc<watch::Sender<bool>>,
     credentials: Credentials,
     environment: Arc<Vec<(String, String)>>,
     /// What the daemon keeps that a sandbox leaves alone or lets through.
@@ -894,6 +897,7 @@ impl Registry {
             detached_budget: DEFAULT_DETACHED_BUDGET,
             #[cfg(unix)]
             terminals: Arc::default(),
+            stopping: Arc::new(watch::Sender::new(false)),
             credentials: Credentials::default(),
             environment: Arc::new(Vec::new()),
             daemon: Arc::default(),
@@ -1714,6 +1718,7 @@ impl Registry {
     ) {
         let registry = self.clone();
         let mut queued = queued;
+        let mut stopping = self.stopping.subscribe();
         tokio::spawn(async move {
             let result = async {
                 let slot = match queued.permit.take() {
@@ -1733,11 +1738,23 @@ impl Registry {
                 )
                 .await?;
                 Ok(registry.shell_outcome(stdout, stderr, exit, elapsed))
+            };
+            // A command still running as the daemon stops is dropped, which
+            // kills its group; its handle then reports the process lost, as
+            // the next start would.
+            tokio::select! {
+                result = result => {
+                    let _ = done.send(result);
+                }
+                _ = stopping.wait_for(|stopping| *stopping) => {}
             }
-            .await;
-            let _ = done.send(result);
         });
     }
+
+    /// End every background command still running, as the daemon stops.
+    pub fn end_background(&self) {
+        self.stopping.send_replace(true);
+    }
 }
 
 /// Number lines from a 1-based offset within the page budget. A single line

tests/test_wait.py

@@ -108,6 +108,24 @@ def background(request, command):
             self.assertLess(process.memory_info().rss - baseline, 8 * 1024 * 1024, phase)
         self.assertEqual(client.request('wait', handles=[pending], timeout_ms=0)['result']['pending'], [pending])
 
+    def test_a_clean_stop_ends_a_running_background_command_and_records_it_lost(self):
+        client = self.client('echo,shell,wait')
+        client.request('create', bot='Bob', workspace=str(self.path))
+        turn = client.request('submit', bot='Bob', request_id='bg',
+                              prompt='bg:while [ ! -f finish ]; do sleep .01; done')['result']['turn']
+        self.assertEqual(client.finished(turn)['data']['status'], 'completed')
+        handle = self.tool_output(client, 'Bob', 'bg-1')['handle']
+        # The stop ends the command and records its end itself, instead of
+        # leaving the row running for the next start to find.
+        client.close()
+        with sqlite3.connect(self.path / 'state.sqlite') as db:
+            rows = db.execute('SELECT status,result FROM processes').fetchall()
+        self.assertEqual([(status, json.loads(result)['error']) for status, result in rows],
+                         [('finished', 'process_lost')])
+        again = self.client('echo,shell,wait')
+        result = again.request('wait', handles=[handle])['result']['results'][handle]
+        self.assertEqual(result['error'], 'process_lost')
+
     def test_old_turn_outcomes_use_indexed_event_lookups(self):
         client = self.client()
         client.request('create', bot='Bob', workspace=str(self.path))
@@ -759,9 +777,24 @@ def test_a_fatal_storage_failure_still_answers_admissions_queued_before_it(self)
         client = self.client('echo,shell,wait')
         for bot in ('Bob', 'Carol'):
             client.request('create', bot=bot, workspace=str(self.path))
+        # The service looks for ready turns once as it starts and again as
+        # each turn ends, on its own loop and through the store. A look still
+        # to come when the lock below is taken would hold that loop, so it
+        # would read the admissions only after the failure: wait for each.
+        def looks(at_least):
+            deadline = time.monotonic() + 10
+            while True:
+                operations = client.request('stats')['result']['store']['operations']
+                count = operations.get('next_ready', {}).get('count', 0)
+                if count >= at_least:
+                    return count
+                self.assertLess(time.monotonic(), deadline)
+                time.sleep(.01)
+        before = looks(1)
         started = client.request('submit', bot='Bob', request_id='bg',
                                  prompt="bg:while [ ! -f release ]; do sleep .01; done; printf done > marker")['result']['turn']
         self.assertEqual(client.finished(started)['data']['status'], 'completed')
+        looks(before + 1)
         with sqlite3.connect(database) as db:
             db.execute("CREATE TRIGGER reject_completion BEFORE UPDATE OF result ON processes BEGIN "
                        "SELECT RAISE(ABORT, 'synthetic write failure'); END")