diff --git a/crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/grok_image_websocket.rs b/crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/grok_image_websocket.rs index 5978f718..74bee910 100644 --- a/crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/grok_image_websocket.rs +++ b/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 { @@ -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(()) } } diff --git a/crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/messages/websocket_server_message.rs b/crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/messages/websocket_server_message.rs index 43f11afc..a12f6031 100644 --- a/crates/api_clients/grok_consumer_client/src/endpoint_bindings/image_websocket/messages/websocket_server_message.rs +++ b/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, + + /// Human-readable message, e.g. "Image rate limit exceeded". + pub err_msg: Option, + + /// UUID of the request this error is a response to. + pub request_id: Option, +} + +/// 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"); diff --git a/crates/api_clients/grok_consumer_client/src/error/grok_generic_api_error.rs b/crates/api_clients/grok_consumer_client/src/error/grok_generic_api_error.rs index a56d0f3a..4c2880ae 100644 --- a/crates/api_clients/grok_consumer_client/src/error/grok_generic_api_error.rs +++ b/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), } diff --git a/crates/api_clients/grok_consumer_client/src/error/grok_specific_api_error.rs b/crates/api_clients/grok_consumer_client/src/error/grok_specific_api_error.rs index e0fd7e40..c2ac1d7f 100644 --- a/crates/api_clients/grok_consumer_client/src/error/grok_specific_api_error.rs +++ b/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."),