this does it!

crates/api_clients/grok_client/src/requests/listen_websocket/clonable_websocket.rs

@@ -2,9 +2,10 @@ use crate::error::grok_client_error::GrokClientError;
 use crate::error::grok_error::GrokError;
 use futures::stream::FusedStream;
 use futures::{Stream, TryStreamExt};
-use log::warn;
+use log::{info, warn};
 use serde::Serialize;
 use std::sync::{Arc, RwLock};
+use std::time::Duration;
 use wreq::ws::message::Message;
 use wreq::ws::WebSocket;
 
@@ -51,6 +52,7 @@ impl ClonableWebsocket {
     }
   }
 
+  /// NB: This will block forever if there is no item.
   pub async fn try_next(&self) -> Result<Option<Message>, GrokClientError> {
     match self.websocket.write() {
       Err(_) => Err(GrokClientError::WebsocketLockError),
@@ -61,6 +63,26 @@ impl ClonableWebsocket {
     }
   }
 
+  pub async fn try_next_timeout(&self, duration: Duration) -> Result<Option<Message>, GrokClientError> {
+    match self.websocket.write() {
+      Err(_) => Err(GrokClientError::WebsocketLockError),
+      Ok(mut websocket) => {
+        let result = tokio::time::timeout(duration, websocket.try_next()).await;
+
+        match result {
+          Err(elapsed) => {
+            info!("Websocket try_next() elapsed: {:?}", elapsed);
+            Ok(None) // Timeout elapsed.
+          }
+          Ok(inner) => {
+            let maybe_message = inner.map_err(GrokClientError::WebsocketReadError)?;
+            Ok(maybe_message)
+          }
+        }
+      },
+    }
+  }
+
   pub fn size_hint(&self) -> Result<(usize, Option<usize>), GrokClientError> {
     match self.websocket.read() {
       Err(_) => Err(GrokClientError::WebsocketLockError),

crates/api_clients/grok_client/src/requests/listen_websocket/prompt_websocket_image.rs

@@ -32,7 +32,7 @@ mod tests {
   #[tokio::test]
   #[ignore] // manually test
   async fn prompt() -> AnyhowResult<()> {
-    setup_test_logging(LevelFilter::Trace);
+    //setup_test_logging(LevelFilter::Trace);
     //let cookies = SESSION_COOKIES_WITHOUT_CF_CLEARANCE;
 
     let cookies = get_test_cookies()?;
@@ -56,7 +56,9 @@ mod tests {
     std::io::stdout().flush()?;
 
     let mut count = 0;
-    
+
+    println!("Polling.");
+
     //while let Some(message) = websocket.try_next().await? {
     //  println!("[{count}] Received websocket message: {:?}", message);
     //  count = count + 1;
@@ -66,23 +68,37 @@ mod tests {
     //  }
     //}
 
-    println!("Polling.");
-    
+    //loop {
+    //  count = count + 1;
+    //  let bounds = websocket.size_hint()?;
+    //  println!("Websocket size hint: {:?}", bounds);
+    //  tokio::time::sleep(Duration::from_millis(1000)).await;
+    //  if count > 30 {
+    //    break;
+    //  }
+    //}
+
     loop {
-      count = count + 1;
-      
-      let bounds = websocket.size_hint()?;
-      
-      println!("Websocket size hint: {:?}", bounds);
+      let maybe_message =
+          websocket.try_next_timeout(Duration::from_millis(1000)).await?;
 
+      match maybe_message {
+        None => {
+          println!("No message received within timeout.");
+          count = count + 1;
+        }
+        Some(message) => {
+          println!("[{count}] Received websocket message!");//: {:?}", message);
+          count = 0;
+        }
+      }
 
-      tokio::time::sleep(Duration::from_millis(1000)).await;
-
-      if count > 30 {
+      if count > 5 {
+        println!("No messages after 5 seconds");
         break;
       }
     }
-    
+
     println!("Done polling.");
 
     log::logger().flush();