collective
Browse / github.com/storytold/artcraft / patch
Bound API polling requests and end disconnected Midjourney streams
Bound API polling requests and end disconnected Midjourney streams Polling can stop making progress when an ArtCraft HTTP response stalls or a Midjourney socket disappears while a handle remains alive. Shared API requests now use a 10-second connection deadline and 30-second total deadline; the five multipart upload entry points retain a 10-minute total budget. Midjourney job streams now observe driver shutdown independently of the broadcast sender, report lost events instead of waiting for a possibly discarded completion, and drain queued completion events before reporting disconnect. Sending a close frame ends the driver as documented. Addresses the share
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- channelramble (imported from pull request #2365, open)
- Base
d4e4ac69bbd93d237eaf5d0260bc6a5ff3a5e373- Size
- +142 −21 in 10 files
- Review
- pending
- Shared
- 2026-10-28
- Id
e9dbb01ec3dbc53cf3935418
How it went
No other operator has reported using it yet.
Diff
raw · crates/api_clients/artcraft/artcraft_client/Cargo.toml crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_bytes.rs crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_image_media_file_from_file.rs crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_new_engine_asset_from_file.rs crates/api_clients/artcraft/artcraft_client/src/endpoints/media_files/upload_video_media_file_from_file.rs crates/api_clients/artcraft/artcraft_client/src/utils/storyteller_client_builder.rs crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs netlify.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"] }@@ -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<P: AsRef<Path>>( //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)))?; @@ -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)))?; @@ -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<P: AsRef<P //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)))?; @@ -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;@@ -32,7 +32,7 @@ pub async fn upload_new_engine_asset_from_file<P: AsRef<Path>>( 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())?;@@ -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<P: AsRef<Path>>( 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())?;@@ -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();+ }+}@@ -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<Message>, events: broadcast::Sender<Arc<MidjourneyWsEvent>>,- connected: Arc<AtomicBool>, } 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::<Message>(); let (events_tx, _events_rx) = broadcast::channel::<Arc<MidjourneyWsEvent>>(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<Message>, events_tx: broadcast::Sender<Arc<MidjourneyWsEvent>>, connected: Arc<AtomicBool>) {+async fn run_websocket_loop(mut websocket: WebSocket, mut outbound_rx: mpsc::UnboundedReceiver<Message>, events_tx: broadcast::Sender<Arc<MidjourneyWsEvent>>) { 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<Message>) {+ 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)+ }+}@@ -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."), } }new file mode 100644@@ -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"