diff --git a/crates/api_clients/artcraft/artcraft_client/Cargo.toml b/crates/api_clients/artcraft/artcraft_client/Cargo.toml --- a/crates/api_clients/artcraft/artcraft_client/Cargo.toml +++ b/crates/api_clients/artcraft/artcraft_client/Cargo.toml @@ -48,4 +48,4 @@ uuid = { version = "1.8.0", features = ["v4"] } url-builder = "0.1.1" [dev-dependencies] -tokio.workspace = true +tokio = { workspace = true, features = ["rt", "time"] } diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs --- a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs @@ -4,7 +4,7 @@ use crate::error::client_error::ClientError; use crate::error::storyteller_error::StorytellerError; use crate::utils::api_host::ApiHost; use crate::utils::constants::APPLICATION_JSON; -use crate::utils::storyteller_client_builder::storyteller_client_builder; +use crate::utils::storyteller_client_builder::storyteller_upload_client_builder; use crate::utils::filter_bad_response::filter_bad_response; use crate::utils::http_get_anonymous::http_get_anonymous; use chrono::{DateTime, Utc}; @@ -56,7 +56,7 @@ pub async fn legacy_upload_media_file_from_file>( //let mut file = File::open(path)?; - let client = storyteller_client_builder(args.api_host) + let client = storyteller_upload_client_builder(args.api_host) .build() .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_bytes.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_bytes.rs --- a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_bytes.rs +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_bytes.rs @@ -4,7 +4,7 @@ use crate::error::client_error::ClientError; use crate::error::storyteller_error::StorytellerError; use crate::utils::api_host::ApiHost; use crate::utils::constants::APPLICATION_JSON; -use crate::utils::storyteller_client_builder::storyteller_client_builder; +use crate::utils::storyteller_client_builder::storyteller_upload_client_builder; use crate::utils::filter_bad_response::filter_bad_response; use crate::utils::http_get_anonymous::http_get_anonymous; use enums::common::generation_provider::GenerationProvider; @@ -72,7 +72,7 @@ pub async fn upload_image_media_file_from_bytes( debug!("Requesting {:?}", &url); - let client = storyteller_client_builder(args.api_host) + let client = storyteller_upload_client_builder(args.api_host) .build() .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_file.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_file.rs --- a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_file.rs +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_file.rs @@ -4,7 +4,7 @@ use crate::error::client_error::ClientError; use crate::error::storyteller_error::StorytellerError; use crate::utils::api_host::ApiHost; use crate::utils::constants::APPLICATION_JSON; -use crate::utils::storyteller_client_builder::storyteller_client_builder; +use crate::utils::storyteller_client_builder::storyteller_upload_client_builder; use crate::utils::filter_bad_response::filter_bad_response; use crate::utils::http_get_anonymous::http_get_anonymous; use chrono::{DateTime, Utc}; @@ -65,7 +65,7 @@ pub async fn upload_image_media_file_from_file_with_idempotency_token>( debug!("Requesting {:?}", &url); - let client = storyteller_client_builder(api_host) + let client = storyteller_upload_client_builder(api_host) .build()?; let file_bytes = std::fs::read(path.as_ref())?; diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_video_media_file_from_file.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_video_media_file_from_file.rs --- a/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_video_media_file_from_file.rs +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_video_media_file_from_file.rs @@ -2,7 +2,7 @@ use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; use crate::error::api_error::ApiError; use crate::utils::api_host::ApiHost; use crate::utils::constants::APPLICATION_JSON; -use crate::utils::storyteller_client_builder::storyteller_client_builder; +use crate::utils::storyteller_client_builder::storyteller_upload_client_builder; use crate::utils::filter_bad_response::filter_bad_response; use crate::utils::http_get_anonymous::http_get_anonymous; use enums::common::generation_provider::GenerationProvider; @@ -58,7 +58,7 @@ pub async fn upload_video_media_file_from_file>( debug!("Requesting {:?}", &url); - let client = storyteller_client_builder(args.api_host) + let client = storyteller_upload_client_builder(args.api_host) .build()?; let file_bytes = std::fs::read(args.path.as_ref())?; diff --git a/crates/api_clients/artcraft/artcraft_client/src/utils/storyteller_client_builder.rs b/crates/api_clients/artcraft/artcraft_client/src/utils/storyteller_client_builder.rs --- a/crates/api_clients/artcraft/artcraft_client/src/utils/storyteller_client_builder.rs +++ b/crates/api_clients/artcraft/artcraft_client/src/utils/storyteller_client_builder.rs @@ -1,9 +1,15 @@ +use std::time::Duration; + use artcraft_client_identity::user_agents::ARTCRAFT_DESKTOP_USER_AGENT; use reqwest::header::{HeaderMap, HeaderValue, ORIGIN}; use reqwest::{Client, ClientBuilder}; use crate::utils::api_host::ApiHost; +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const API_TIMEOUT: Duration = Duration::from_secs(30); +const UPLOAD_TIMEOUT: Duration = Duration::from_secs(600); + /// Every request to our API is built from this so the desktop client identifies itself /// consistently (`User-Agent` and `Origin`). Only use it for Storyteller / ArtCraft APIs. pub(crate) fn storyteller_client_builder(api_host: &ApiHost) -> ClientBuilder { @@ -13,5 +19,53 @@ pub(crate) fn storyteller_client_builder(api_host: &ApiHost) -> ClientBuilder { Client::builder() .user_agent(ARTCRAFT_DESKTOP_USER_AGENT) .default_headers(headers) + .connect_timeout(CONNECT_TIMEOUT) + .timeout(API_TIMEOUT) .gzip(true) } + +/// Multipart media uploads can take longer than polling and small JSON calls. +pub(crate) fn storyteller_upload_client_builder(api_host: &ApiHost) -> ClientBuilder { + storyteller_client_builder(api_host).timeout(UPLOAD_TIMEOUT) +} + +#[cfg(test)] +mod timeout_fixture_tests { + use super::*; + use std::io::{Read, Write}; + use std::net::TcpListener; + + #[tokio::test] + async fn stalled_headers_time_out_and_the_next_poll_succeeds() { + check_stalled_response(false).await; + } + + #[tokio::test] + async fn stalled_body_times_out_and_the_next_poll_succeeds() { + check_stalled_response(true).await; + } + + async fn check_stalled_response(send_headers: bool) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let server = std::thread::spawn(move || { + let (mut stalled, _) = listener.accept().unwrap(); + let mut buffer = [0; 4096]; + stalled.read(&mut buffer).unwrap(); + if send_headers { + stalled.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 20\r\nConnection: close\r\n\r\n{").unwrap(); + } + let (mut healthy, _) = listener.accept().unwrap(); + healthy.read(&mut buffer).unwrap(); + healthy.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}").unwrap(); + }); + let client = storyteller_client_builder(&ApiHost::Localhost { port: address.port() as u32 }).no_proxy().build().unwrap(); + let url = format!("http://{address}/jobs"); + let result = tokio::time::timeout(API_TIMEOUT + Duration::from_secs(5), async { + client.get(&url).send().await?.text().await + }).await.expect("the configured deadline must end the stalled request"); + assert!(result.unwrap_err().is_timeout()); + assert_eq!(client.get(&url).send().await.unwrap().text().await.unwrap(), "{}"); + server.join().unwrap(); + } +} diff --git a/crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs b/crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs --- a/crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs +++ b/crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs @@ -3,7 +3,6 @@ use crate::credentials::midjourney_user_id::MidjourneyUserId; use crate::error::midjourney_client_error::MidjourneyClientError; use log::{debug, warn}; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; use tokio::sync::{broadcast, mpsc}; use wreq::ws::WebSocket; use wreq::ws::message::Message; @@ -46,7 +45,6 @@ struct WebSocketInner { user_id: MidjourneyUserId, outbound: mpsc::UnboundedSender, events: broadcast::Sender>, - connected: Arc, } impl MidjourneyWebSocket { @@ -56,11 +54,10 @@ impl MidjourneyWebSocket { pub(crate) fn spawn(user_id: MidjourneyUserId, websocket: WebSocket) -> Self { let (outbound_tx, outbound_rx) = mpsc::unbounded_channel::(); let (events_tx, _events_rx) = broadcast::channel::>(EVENT_CHANNEL_CAPACITY); - let connected = Arc::new(AtomicBool::new(true)); - tokio::spawn(run_websocket_loop(websocket, outbound_rx, events_tx.clone(), connected.clone())); + tokio::spawn(run_websocket_loop(websocket, outbound_rx, events_tx.clone())); - Self { inner: Arc::new(WebSocketInner { user_id, outbound: outbound_tx, events: events_tx, connected }) } + Self { inner: Arc::new(WebSocketInner { user_id, outbound: outbound_tx, events: events_tx }) } } /// The user this connection is bound to. @@ -70,7 +67,7 @@ impl MidjourneyWebSocket { /// Whether the background task is still running. pub fn is_connected(&self) -> bool { - self.inner.connected.load(Ordering::SeqCst) + !self.inner.outbound.is_closed() } /// A new receiver over the decoded event stream. Each call yields an @@ -113,7 +110,14 @@ impl MidjourneyWebSocket { self.subscribe_to_job(job_id)?; loop { - match receiver.recv().await { + // Drain queued terminal events before reporting that the driver stopped. The + // event sender lives in this handle too, so its channel alone cannot signal EOF. + let event = tokio::select! { + biased; + event = receiver.recv() => event, + _ = self.inner.outbound.closed() => return Err(MidjourneyClientError::WebSocketClosed), + }; + match event { Ok(event) => { if event.job_id() != Some(job_id) { continue; @@ -124,7 +128,7 @@ impl MidjourneyWebSocket { } }, Err(broadcast::error::RecvError::Lagged(skipped)) => { - warn!("Midjourney websocket consumer lagged, skipped {} events", skipped); + return Err(MidjourneyClientError::WebSocketEventsLost(skipped)); }, Err(broadcast::error::RecvError::Closed) => { return Err(MidjourneyClientError::WebSocketClosed); @@ -146,16 +150,18 @@ impl MidjourneyWebSocket { /// Owns the raw socket: forwards outbound command frames and decodes inbound /// CBOR frames into events. Ends when the socket closes or all handles drop. -async fn run_websocket_loop(mut websocket: WebSocket, mut outbound_rx: mpsc::UnboundedReceiver, events_tx: broadcast::Sender>, connected: Arc) { +async fn run_websocket_loop(mut websocket: WebSocket, mut outbound_rx: mpsc::UnboundedReceiver, events_tx: broadcast::Sender>) { loop { tokio::select! { maybe_outbound = outbound_rx.recv() => { match maybe_outbound { Some(message) => { + let closing = matches!(&message, Message::Close(_)); if let Err(err) = websocket.send(message).await { warn!("Midjourney websocket send failed: {}", err); break; } + if closing { break; } } // All handles dropped; nothing more will be sent. None => break, @@ -185,6 +191,56 @@ async fn run_websocket_loop(mut websocket: WebSocket, mut outbound_rx: mpsc::Unb } } - connected.store(false, Ordering::SeqCst); debug!("Midjourney websocket loop ended."); } + +#[cfg(test)] +mod lifecycle_tests { + use super::*; + use std::time::Duration; + + #[tokio::test] + async fn disconnect_ends_stream_while_the_handle_is_retained() { + let (socket, mut outbound) = fixture(4); + let stream = socket.stream_job("job", |_| {}); + let disconnect = async { outbound.recv().await.unwrap(); drop(outbound); }; + let (result, _) = tokio::time::timeout(Duration::from_secs(1), async { tokio::join!(stream, disconnect) }).await.unwrap(); + assert!(matches!(result, Err(MidjourneyClientError::WebSocketClosed))); + assert!(!socket.is_connected()); + } + + #[tokio::test] + async fn losing_a_terminal_event_reports_an_error() { + let (socket, mut outbound) = fixture(1); + let stream = socket.stream_job("job", |_| {}); + let overflow = async { + outbound.recv().await.unwrap(); + socket.inner.events.send(Arc::new(MidjourneyWsEvent::Completed { job_id: "job".into(), percentage_complete: 100.0 })).unwrap(); + socket.inner.events.send(Arc::new(MidjourneyWsEvent::Unknown { maybe_kind: None })).unwrap(); + }; + let (result, _) = tokio::time::timeout(Duration::from_secs(1), async { tokio::join!(stream, overflow) }).await.unwrap(); + assert!(matches!(result, Err(MidjourneyClientError::WebSocketEventsLost(1)))); + } + + #[tokio::test] + async fn queued_completion_wins_over_disconnect() { + let (socket, mut outbound) = fixture(4); + let mut completed = false; + let stream = socket.stream_job("job", |_| completed = true); + let finish = async { + outbound.recv().await.unwrap(); + socket.inner.events.send(Arc::new(MidjourneyWsEvent::Completed { job_id: "job".into(), percentage_complete: 100.0 })).unwrap(); + drop(outbound); + }; + let (result, _) = tokio::join!(stream, finish); + result.unwrap(); + assert!(completed); + } + + fn fixture(capacity: usize) -> (MidjourneyWebSocket, mpsc::UnboundedReceiver) { + let (outbound, receiver) = mpsc::unbounded_channel(); + let (events, _) = broadcast::channel(capacity); + let socket = MidjourneyWebSocket { inner: Arc::new(WebSocketInner { user_id: MidjourneyUserId::from_str("test"), outbound, events }) }; + (socket, receiver) + } +} diff --git a/crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs b/crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs --- a/crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs +++ b/crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs @@ -34,6 +34,9 @@ pub enum MidjourneyClientError { /// Tried to send on a websocket whose background task has stopped. WebSocketClosed, + /// Progress was lost, possibly including the terminal event. Reconcile over HTTP. + WebSocketEventsLost(u64), + /// The websocket opened but the `subscribe_to_user` handshake did not /// complete within the allotted time. WebSocketHandshakeTimeout, @@ -54,6 +57,7 @@ impl Display for MidjourneyClientError { Self::SerializationError(err) => write!(f, "Request serialization error: {}", err), Self::WebSocketUpgrade(msg) => write!(f, "WebSocket upgrade failed: {}", msg), Self::WebSocketClosed => write!(f, "The websocket connection is closed."), + Self::WebSocketEventsLost(count) => write!(f, "The websocket lost {count} events; refresh the job status."), Self::WebSocketHandshakeTimeout => write!(f, "Timed out waiting for the websocket user handshake."), } } diff --git a/netlify.toml b/netlify.toml new file mode 100644 --- /dev/null +++ b/netlify.toml @@ -0,0 +1,7 @@ +# The `storyteller-docs` Netlify site still builds this repository, but its +# build command (./script/netlify_build_docs.sh) and the docs themselves now +# live in storytold/artcraft-services, so every deploy preview here fails. +# Skip those builds until the Netlify site is repointed or disconnected, then +# delete this file. +[build] + ignore = "exit 0"