Give a run one handle, its own identity, and pluggable hooks

.makefiles/rust.mk

@@ -57,6 +57,10 @@ coverage: $(DOCKER_ENV) $(REPORT_DIR) $(GRCOV_BIN) ## Run the local test suite w
 nextest: $(DOCKER_ENV) $(NEXTEST_BIN) $(REPORT_DIR)
 	$(RUN) cargo nextest run --workspace --all-targets --features integration $(if $(IS_CI),-P ci,)
 
+.PHONY:test-loom
+test-loom: $(DOCKER_ENV) ## Run the loom interleaving models (not part of `test`; they explore, so they are slow)
+	$(RUN) env RUSTFLAGS="--cfg aura_loom" cargo test -p aura-web-server --lib loom_
+
 .PHONY:lint-rust
 lint-rust: | $(DOCKER_ENV) $(REPORT_DIR)  ## lint rust code via clippy
 	$(RUN) cargo clippy $(if $(IS_CI),-q,) --all-targets --all-features $(if $(IS_CI),--message-format=json,) -- -D warnings $(if $(IS_CI),> $(REPORT_DIR)/clippy.json,)

crates/aura-test-utils/src/mock_agent.rs

@@ -6,11 +6,10 @@ use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use async_trait::async_trait;
+use aura::streaming::AgentRun;
 use aura::{Message, StreamError, StreamItem, StreamingAgent, UsageState};
 use aura::{StreamedAssistantContent, StreamedUserContent, ToolCall, ToolResult};
 use futures::stream::{self, BoxStream, StreamExt};
-use tokio::sync::watch;
-use tokio_util::sync::CancellationToken;
 
 type StartHook = Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
 type EffectHook = Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
@@ -107,7 +106,7 @@ impl MockAgent {
     }
 
     /// Runs the given steps in order, then ends the stream. The script is
-    /// consumed by the first `stream`/`stream_with_timeout` call; a second call
+    /// consumed by the first `stream` call; a second call
     /// on the same agent yields an empty stream.
     pub fn scripted(steps: Vec<Step>) -> Self {
         Self {
@@ -116,7 +115,7 @@ impl MockAgent {
         }
     }
 
-    /// Awaited before either entry point produces its stream, and passed that
+    /// Awaited before `stream` produces its stream, and passed that
     /// call's `request_id`.
     pub fn on_stream_start<F, Fut>(mut self, hook: F) -> Self
     where
@@ -175,26 +174,18 @@ impl StreamingAgent for MockAgent {
         &self,
         _query: &str,
         _chat_history: Vec<Message>,
-        _cancel_token: CancellationToken,
+        options: aura::streaming::RunOptions,
         request_id: &str,
-    ) -> Result<BoxStream<'static, Result<StreamItem, StreamError>>, StreamError> {
-        Ok(self.start(request_id).await)
-    }
-
-    async fn stream_with_timeout(
-        &self,
-        _query: &str,
-        _chat_history: Vec<Message>,
-        _timeout: Duration,
-        request_id: &str,
-    ) -> (
-        BoxStream<'static, Result<StreamItem, StreamError>>,
-        watch::Sender<bool>,
-        UsageState,
-    ) {
+    ) -> AgentRun {
         let stream = self.start(request_id).await;
-        let (cancel_tx, _cancel_rx) = watch::channel(false);
-        (stream, cancel_tx, UsageState::new())
+        // Carries a caller-supplied token so `cancel_token()` returns the one the
+        // caller named. The scripts do not race it, so cancelling does not end a
+        // mock stream.
+        AgentRun::new(
+            stream,
+            options.into_parts().1.unwrap_or_default(),
+            UsageState::new(),
+        )
     }
 
     async fn cancel_and_close_mcp(&self, _request_id: &str, _reason: &str) -> usize {
@@ -211,68 +202,55 @@ mod tests {
     async fn a_pending_agent_never_yields() {
         let agent = MockAgent::pending();
         let mut stream = agent
-            .stream("q", vec![], CancellationToken::new(), "req_1")
+            .stream("q", vec![], aura::streaming::RunOptions::default(), "req_1")
             .await
-            .expect("mock stream should start");
+            .into_events();
         assert!(
             tokio::time::timeout(Duration::from_millis(50), stream.next())
                 .await
                 .is_err()
         );
     }
 
-    #[tokio::test]
-    async fn the_start_hook_runs_before_both_entry_points_stream() {
-        for entry_point in ["stream", "stream_with_timeout"] {
-            let ran = Arc::new(AtomicBool::new(false));
-            let flag = Arc::clone(&ran);
-            let agent = MockAgent::pending().on_stream_start(move |request_id| {
-                let flag = Arc::clone(&flag);
-                async move {
-                    assert_eq!(request_id, "req_1");
-                    flag.store(true, Ordering::SeqCst);
-                }
-            });
-
-            if entry_point == "stream" {
-                let _ = agent
-                    .stream("q", vec![], CancellationToken::new(), "req_1")
-                    .await
-                    .expect("mock stream should start");
-            } else {
-                let _ = agent
-                    .stream_with_timeout("q", vec![], Duration::from_secs(1), "req_1")
-                    .await;
+    /// `into_events` wraps the stream, so the order is asserted rather than the
+    /// count — a wrapper that buffered or reordered would keep the count.
+    #[tokio::test(start_paused = true)]
+    async fn a_yielding_agent_produces_its_items_then_ends() {
+        let agent = MockAgent::yielding(vec![items::text("hello "), items::text("world")]);
+        let mut stream = agent
+            .stream("q", vec![], aura::streaming::RunOptions::default(), "req_1")
+            .await
+            .into_events();
+
+        let mut texts = Vec::new();
+        while let Some(item) = stream.next().await {
+            if let Ok(aura::StreamItem::StreamAssistantItem(StreamedAssistantContent::Text(t))) =
+                item
+            {
+                texts.push(t);
             }
-
-            assert!(
-                ran.load(Ordering::SeqCst),
-                "hook should run for {entry_point}"
-            );
         }
+        assert_eq!(texts, ["hello ", "world"]);
     }
 
-    #[tokio::test(start_paused = true)]
-    async fn a_yielding_agent_produces_its_items_then_ends() {
-        let agent = MockAgent::yielding([items::text("hello "), items::text("world")]);
-        let stream = agent
-            .stream("q", vec![], CancellationToken::new(), "req_1")
+    #[tokio::test]
+    async fn the_start_hook_runs_before_the_stream() {
+        let ran = Arc::new(AtomicBool::new(false));
+        let flag = Arc::clone(&ran);
+        let agent = MockAgent::pending().on_stream_start(move |request_id| {
+            let flag = Arc::clone(&flag);
+            async move {
+                assert_eq!(request_id, "req_1");
+                flag.store(true, Ordering::SeqCst);
+            }
+        });
+
+        let _ = agent
+            .stream("q", vec![], aura::streaming::RunOptions::default(), "req_1")
             .await
-            .expect("mock stream should start");
-
-        let texts: Vec<String> = stream
-            .filter_map(|item| async move {
-                match item {
-                    Ok(StreamItem::StreamAssistantItem(StreamedAssistantContent::Text(t))) => {
-                        Some(t)
-                    }
-                    _ => None,
-                }
-            })
-            .collect()
-            .await;
+            .into_events();
 
-        assert_eq!(texts, vec!["hello ", "world"]);
+        assert!(ran.load(Ordering::SeqCst), "hook should run");
     }
 
     #[tokio::test(start_paused = true)]
@@ -290,9 +268,14 @@ mod tests {
         ]);
 
         let stream = agent
-            .stream("q", vec![], CancellationToken::new(), "req_42")
+            .stream(
+                "q",
+                vec![],
+                aura::streaming::RunOptions::default(),
+                "req_42",
+            )
             .await
-            .expect("mock stream should start");
+            .into_events();
         let items: Vec<_> = stream.collect().await;
 
         assert_eq!(order.lock().expect("order lock").as_slice(), ["req_42"]);
@@ -304,15 +287,15 @@ mod tests {
         let agent = MockAgent::yielding([items::text("once")]);
 
         let first: Vec<_> = agent
-            .stream("q", vec![], CancellationToken::new(), "req_1")
+            .stream("q", vec![], aura::streaming::RunOptions::default(), "req_1")
             .await
-            .expect("mock stream should start")
+            .into_events()
             .collect()
             .await;
         let second: Vec<_> = agent
-            .stream("q", vec![], CancellationToken::new(), "req_1")
+            .stream("q", vec![], aura::streaming::RunOptions::default(), "req_1")
             .await
-            .expect("mock stream should start")
+            .into_events()
             .collect()
             .await;
 
@@ -337,9 +320,9 @@ mod tests {
         ]);
 
         let mut stream = agent
-            .stream("q", vec![], CancellationToken::new(), "req_1")
+            .stream("q", vec![], aura::streaming::RunOptions::default(), "req_1")
             .await
-            .expect("mock stream should start");
+            .into_events();
         while stream.next().await.is_some() {
             seen.fetch_add(1, Ordering::SeqCst);
         }

crates/aura-web-server/Cargo.toml

@@ -98,3 +98,13 @@ aura-events = { path = "../aura-events" }
 tower = "0.5"
 http-body-util = "0.1"
 tempfile = "3"
+
+# Exhaustively explores interleavings of the task-cancel map, which a stress
+# test cannot reach: the window is one mutex release and reacquire wide.
+[lints.rust]
+unexpected_cfgs = { level = "warn", check-cfg = ['cfg(aura_loom)'] }
+
+# Not a dev-dependency: the lib itself swaps in loom's mutex under this cfg,
+# which is never set by a normal build.
+[target.'cfg(aura_loom)'.dependencies]
+loom = "0.7"