handle error

crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/grok_image_websocket.rs

@@ -1,7 +1,8 @@
 use crate::client::browser_user_agents::FIREFOX_143_MAC_USER_AGENT;
 use crate::datatypes::api::request_id::RequestId;
 use crate::endpoint_bindings::image_websocket::messages::websocket_client_message::{ClientMessageAspectRatio, WebsocketClientMessage};
-use crate::endpoint_bindings::image_websocket::messages::websocket_server_message::WebsocketServerMessage;
+use crate::endpoint_bindings::image_websocket::messages::websocket_server_message::{ErrorMessage, ERR_CODE_RATE_LIMIT_EXCEEDED, WebsocketServerMessage};
+use crate::error::grok_specific_api_error::GrokSpecificApiError;
 use crate::error::grok_client_error::GrokClientError;
 use crate::error::grok_error::GrokError;
 use crate::error::grok_generic_api_error::GrokGenericApiError;
@@ -167,6 +168,12 @@ impl GrokImageWebsocket {
         continue;
       };
 
+      // An error frame (e.g. out-of-quota) means no images are coming; fail
+      // fast with a typed error carrying the raw frame.
+      if let WebsocketServerMessage::Error(error) = &message {
+        return Err(error_frame_to_grok_error(error, &text));
+      }
+
       if let Some(image) = completed_image(&message) {
         images.push(image);
       }
@@ -192,6 +199,20 @@ impl GrokImageWebsocket {
   }
 }
 
+/// Map a server error frame (plus its raw body) to a typed error. Quota /
+/// rate-limit exhaustion gets its own [`GrokSpecificApiError`] variant; any
+/// other code is surfaced generically. Both carry the full raw frame.
+fn error_frame_to_grok_error(error: &ErrorMessage, raw_frame: &str) -> GrokError {
+  match error.err_code.as_deref() {
+    Some(ERR_CODE_RATE_LIMIT_EXCEEDED) => {
+      GrokSpecificApiError::ImageRateLimitExceeded { body: raw_frame.to_string() }.into()
+    }
+    _ => {
+      GrokGenericApiError::UnexpectedWebsocketErrorFrame { body: raw_frame.to_string() }.into()
+    }
+  }
+}
+
 /// The completed-image view of a server message, if it is a finished image
 /// (100%) that carries a URL.
 fn completed_image(message: &WebsocketServerMessage) -> Option<CompletedImage> {
@@ -274,13 +295,22 @@ mod tests {
       ClientMessageAspectRatio::WideThreeByTwo,
     ).await?;
 
-    let images = websocket.collect_images(Duration::from_secs(30)).await?;
-    info!("Collected {} image(s).", images.len());
-    for image in &images {
-      info!("Image: {} ({})", image.url, image.request_id.to_string());
+    match websocket.collect_images(Duration::from_secs(30)).await {
+      Ok(images) => {
+        info!("Collected {} image(s).", images.len());
+        for image in &images {
+          info!("Image: {} ({})", image.url, image.request_id.to_string());
+        }
+        assert!(!images.is_empty(), "expected at least one completed image");
+      }
+      // Out of quota is a valid "the plumbing works" outcome; surface it
+      // clearly rather than failing.
+      Err(GrokError::ApiSpecific(GrokSpecificApiError::ImageRateLimitExceeded { body })) => {
+        info!("Image rate limit / quota exhausted (expected when out of quota): {body}");
+      }
+      Err(err) => return Err(err.into()),
     }
 
-    assert!(!images.is_empty(), "expected at least one completed image");
     Ok(())
   }
 }

crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/messages/websocket_server_message.rs

@@ -12,12 +12,37 @@ pub enum WebsocketServerMessage {
   #[serde(rename = "json")]
   Json(JsonDataMessage),
 
+  /// An error frame, root `type = "error"`. e.g. when the image rate limit /
+  /// quota is exhausted.
+  #[serde(rename = "error")]
+  Error(ErrorMessage),
+
   /// Captures any other JSON messages.
   #[serde(untagged)]
   Unknown(serde_json::Value),
   //Unknown(String),
 }
 
+/// An error the server sends over the socket instead of generating.
+///
+/// e.g. rate limiting:
+/// `{"type":"error","err_code":"rate_limit_exceeded","err_msg":"Image rate limit exceeded","request_id":"..."}`
+#[derive(Deserialize, Clone, Debug)]
+pub struct ErrorMessage {
+  /// Machine-readable code, e.g. "rate_limit_exceeded".
+  pub err_code: Option<String>,
+
+  /// Human-readable message, e.g. "Image rate limit exceeded".
+  pub err_msg: Option<String>,
+
+  /// UUID of the request this error is a response to.
+  pub request_id: Option<String>,
+}
+
+/// The `err_code` on an [`ErrorMessage`] when the account's image quota / rate
+/// limit is exhausted.
+pub const ERR_CODE_RATE_LIMIT_EXCEEDED: &str = "rate_limit_exceeded";
+
 /// Images with a binary blob and URLs.
 /// We may receive several of these for a single prompt.
 #[derive(Deserialize, Clone, Debug)]
@@ -84,13 +109,33 @@ impl WebsocketServerMessage {
 
 #[cfg(test)]
 mod tests {
-  use crate::endpoint_bindings::image_websocket::messages::websocket_server_message::WebsocketServerMessage;
+  use crate::endpoint_bindings::image_websocket::messages::websocket_server_message::{ERR_CODE_RATE_LIMIT_EXCEEDED, WebsocketServerMessage};
 
   fn json_path(file_name: &str) -> String {
     // NB: Cargo runs tests with the crate root as the working directory.
     format!("test_data/websocket_messages/{}", file_name)
   }
 
+  // Real rate-limit frame captured over the socket when the account is out of
+  // image quota.
+  const RATE_LIMIT_FRAME: &str = r#"{"type":"error","err_code":"rate_limit_exceeded","err_msg":"Image rate limit exceeded","request_id":"5c8fdd50-11ab-4e7c-9f64-301902eb6678"}"#;
+
+  #[test]
+  fn test_error_rate_limit() -> anyhow::Result<()> {
+    let message = WebsocketServerMessage::from_json_str(RATE_LIMIT_FRAME)?;
+
+    match message {
+      WebsocketServerMessage::Error(error) => {
+        assert_eq!(error.err_code.as_deref(), Some(ERR_CODE_RATE_LIMIT_EXCEEDED));
+        assert_eq!(error.err_msg.as_deref(), Some("Image rate limit exceeded"));
+        assert!(error.request_id.is_some());
+      }
+      other => panic!("Expected Error message, got {:?}", other),
+    }
+
+    Ok(())
+  }
+
   #[test]
   fn test_image_complete() -> anyhow::Result<()> {
     let path = json_path("create_image_complete_response.json");

crates/api_clients/grok_consumer_client/src/error/grok_generic_api_error.rs

@@ -52,6 +52,12 @@ pub enum GrokGenericApiError {
   /// Our previous upload failed
   UploadFailed,
 
+  /// The image websocket returned an error frame we don't specifically handle.
+  /// Carries the full raw error frame.
+  UnexpectedWebsocketErrorFrame {
+    body: String,
+  },
+
   /// An error upgrading the connection to a websocket.
   WreqWebsocketUpgradeError(wreq::Error),
 
@@ -75,6 +81,7 @@ impl Display for GrokGenericApiError {
       Self::UncategorizedBadResponseWithStatus(status) => write!(f, "Uncategorized with status code: {}", status),
       Self::UncategorizedBadResponseWithStatusAndBody { status_code, body } => write!(f, "Uncategorized bad response: status code {}, body: {}", status_code, body),
       Self::UploadFailed => write!(f, "Upload failed"),
+      Self::UnexpectedWebsocketErrorFrame { body } => write!(f, "Unexpected websocket error frame: {}", body),
       Self::WreqWebsocketUpgradeError(err) => write!(f, "Websocket upgrade error: {}", err),
       Self::WreqError(err) => write!(f, "Wreq client error: {}", err),
     }

crates/api_clients/grok_consumer_client/src/error/grok_specific_api_error.rs

@@ -8,7 +8,14 @@ pub enum GrokSpecificApiError {
   
   /// Too many videos are being generated.
   TooManyVideos,
-  
+
+  /// The image websocket rejected the prompt because the account's image
+  /// generation rate limit / quota is exhausted. Carries the full raw error
+  /// frame the server sent.
+  ImageRateLimitExceeded {
+    body: String,
+  },
+
 
 //  /// Failed to generate a bearer token despite the `200 OK` response.
 //  /// This has been generated by sending old/expired session cookies or empty cookies.
@@ -62,6 +69,7 @@ impl Display for GrokSpecificApiError {
     match self {
       Self::AutomationBlocked => write!(f, "Automation blocked"),
       Self::TooManyVideos => write!(f, "Too many videos are being generated"),
+      Self::ImageRateLimitExceeded { body } => write!(f, "Image rate limit / quota exceeded. Raw frame: {}", body),
 
       //Self::FailedToGenerateBearerToken => write!(f, "Failed to generate bearer token. Probably invalid or expired session cookies."),
       //Self::InvalidJwt => write!(f, "Invalid JWT token. Please refresh the token."),