first pass of merging in warp (doesn't build)

app/src/ai/agent_events/message_hydrator_tests.rs

@@ -1,4 +1,6 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
+use std::time::Duration;
 
 use mockall::predicate::eq;
 
@@ -6,6 +8,7 @@ use super::*;
 use crate::server::server_api::ai::{
     AIClient, AgentRunEvent, MockAIClient, ReadAgentMessageResponse,
 };
+use crate::server::server_api::presigned_upload::HttpStatusError;
 
 fn make_run_event(
     sequence: i64,
@@ -23,6 +26,25 @@ fn make_run_event(
     }
 }
 
+fn make_message_response(message_id: &str) -> ReadAgentMessageResponse {
+    ReadAgentMessageResponse {
+        message_id: message_id.to_string(),
+        sender_run_id: "parent-run".to_string(),
+        subject: "Need a redirect".to_string(),
+        body: "Switch to the failing test first.".to_string(),
+        sent_at: "2026-01-01T00:00:00Z".to_string(),
+        delivered_at: None,
+        read_at: Some("2026-01-01T00:00:01Z".to_string()),
+    }
+}
+
+fn http_status_read_error(status: u16) -> anyhow::Error {
+    anyhow::Error::new(HttpStatusError {
+        status,
+        body: format!("status {status} body"),
+    })
+}
+
 #[tokio::test]
 async fn hydrator_reads_new_message_for_matching_run() {
     let mut ai_client = MockAIClient::new();
@@ -30,17 +52,7 @@ async fn hydrator_reads_new_message_for_matching_run() {
         .expect_read_agent_message()
         .with(eq("msg-123"))
         .times(1)
-        .returning(|_| {
-            Ok(ReadAgentMessageResponse {
-                message_id: "msg-123".to_string(),
-                sender_run_id: "parent-run".to_string(),
-                subject: "Need a redirect".to_string(),
-                body: "Switch to the failing test first.".to_string(),
-                sent_at: "2026-01-01T00:00:00Z".to_string(),
-                delivered_at: None,
-                read_at: Some("2026-01-01T00:00:01Z".to_string()),
-            })
-        });
+        .returning(|_| Ok(make_message_response("msg-123")));
 
     let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
     let hydrator = MessageHydrator::new(ai_client);
@@ -69,3 +81,109 @@ async fn hydrator_ignores_events_for_other_runs() {
         .await
         .is_none());
 }
+
+#[tokio::test]
+async fn read_message_with_timeout_retries_transient_failures_until_success() {
+    let attempts = Arc::new(AtomicUsize::new(0));
+    let attempts_clone = attempts.clone();
+    let mut ai_client = MockAIClient::new();
+    ai_client
+        .expect_read_agent_message()
+        .with(eq("msg-123"))
+        .times(2)
+        .returning(move |_| {
+            let attempt = attempts_clone.fetch_add(1, Ordering::SeqCst);
+            if attempt == 0 {
+                Err(http_status_read_error(404))
+            } else {
+                Ok(make_message_response("msg-123"))
+            }
+        });
+
+    let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
+    let hydrator = MessageHydrator::with_fetch_timing(
+        ai_client,
+        Duration::from_millis(100),
+        Duration::from_millis(5),
+    );
+
+    let message = hydrator.read_message_with_timeout("msg-123").await.unwrap();
+
+    assert_eq!(attempts.load(Ordering::SeqCst), 2);
+    assert_eq!(message.message_id, "msg-123");
+    assert_eq!(message.body, "Switch to the failing test first.");
+}
+
+#[tokio::test]
+async fn read_message_with_timeout_times_out_after_retrying_transient_failures() {
+    let attempts = Arc::new(AtomicUsize::new(0));
+    let attempts_clone = attempts.clone();
+    let mut ai_client = MockAIClient::new();
+    ai_client
+        .expect_read_agent_message()
+        .with(eq("msg-123"))
+        .returning(move |_| {
+            attempts_clone.fetch_add(1, Ordering::SeqCst);
+            Err(http_status_read_error(404))
+        });
+
+    let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
+    let hydrator = MessageHydrator::with_fetch_timing(
+        ai_client,
+        Duration::from_millis(120),
+        Duration::from_millis(20),
+    );
+
+    let err = hydrator
+        .read_message_with_timeout("msg-123")
+        .await
+        .expect_err("expected timeout after transient retries");
+    let err_chain = format!("{err:#}");
+
+    assert!(
+        err_chain.contains("Timed out reading agent message msg-123"),
+        "{err:#}"
+    );
+    assert!(
+        attempts.load(Ordering::SeqCst) >= 2,
+        "expected at least one retry before timeout"
+    );
+}
+
+#[tokio::test]
+async fn read_message_with_timeout_does_not_retry_permanent_http_failures() {
+    let attempts = Arc::new(AtomicUsize::new(0));
+    let attempts_clone = attempts.clone();
+    let mut ai_client = MockAIClient::new();
+    ai_client
+        .expect_read_agent_message()
+        .with(eq("msg-123"))
+        .times(1)
+        .returning(move |_| {
+            attempts_clone.fetch_add(1, Ordering::SeqCst);
+            Err(http_status_read_error(403))
+        });
+
+    let ai_client: Arc<dyn AIClient> = Arc::new(ai_client);
+    let hydrator = MessageHydrator::with_fetch_timing(
+        ai_client,
+        Duration::from_millis(100),
+        Duration::from_millis(5),
+    );
+
+    let err = hydrator
+        .read_message_with_timeout("msg-123")
+        .await
+        .expect_err("expected permanent failure to fail fast");
+    let err_chain = format!("{err:#}");
+
+    assert!(
+        err_chain.contains("HTTP request failed with status 403"),
+        "{err:#}"
+    );
+    assert_eq!(
+        attempts.load(Ordering::SeqCst),
+        1,
+        "permanent 4xx errors should not retry"
+    );
+}