From 584fdb70443e546ef18d9bc9778ec2b343cc87c8 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 14 Sep 2026 12:45:24 -0400 Subject: [PATCH] backport midjourney fixes from ArtCraft-X (#1953) --- Cargo.lock | 62 +- Cargo.toml | 4 + .../upload_image_media_file_from_file.rs | 11 +- .../artcraft/artcraft_router/Cargo.toml | 2 + .../src/api/router_provider.rs | 1 + .../artcraft_router/src/client/mod.rs | 2 + .../src/client/multi_router_client.rs | 7 + .../src/client/multi_router_client_builder.rs | 9 + .../src/client/router_client.rs | 10 + .../src/client/router_midjourney_client.rs | 24 + .../src/errors/artcraft_router_error.rs | 1 + .../src/errors/client_error.rs | 2 + .../src/errors/provider_error.rs | 6 + .../generate_image_request_builder.rs | 2 + .../generate_image/generate_image_response.rs | 6 + .../image_generation_request.rs | 6 + .../providers/midjourney/build.rs | 130 ++++ .../providers/midjourney/cost.rs | 15 + .../providers/midjourney/mod.rs | 3 + .../providers/midjourney/request.rs | 41 ++ .../generate/generate_image/providers/mod.rs | 2 + .../api_clients/midjourney_client/Cargo.toml | 7 +- .../src/client/midjourney_websocket_client.rs | 0 .../midjourney_client/src/client/mod.rs | 2 +- .../src/client/websocket/job_progress.rs | 117 ++++ .../client/websocket/midjourney_websocket.rs | 190 ++++++ .../client/websocket/midjourney_ws_event.rs | 231 +++++++ .../src/client/websocket/mod.rs | 15 + .../websocket/open_midjourney_websocket.rs | 132 ++++ .../cookie_store_has_auth_cookies.rs | 2 +- .../src/endpoints/get_index_page_html.rs | 60 +- .../src/endpoints/imagine.rs | 69 +- .../midjourney_client/src/endpoints/mod.rs | 4 +- .../src/endpoints/prompt_session_log.rs | 108 ++++ .../src/endpoints/refresh_token.rs | 96 +++ .../src/endpoints/storage_list.rs | 55 +- .../src/endpoints/submit_job.rs | 155 ++++- .../src/endpoints/user_connections.rs | 56 +- .../src/error/midjourney_api_error.rs | 21 +- .../src/error/midjourney_client_error.rs | 13 + .../get_midjourney_user_id_from_storage.rs | 36 +- .../src/recipes/get_user_info.rs | 47 +- .../src/recipes/text_to_image.rs | 49 +- .../src/utils/image_downloader_client.rs | 21 +- crates/desktop/artcraft/Cargo.toml | 3 + .../get_app_preferences_command.rs | 5 + .../update_app_preference_command.rs | 60 +- .../estimate_image_cost_command.rs | 19 + .../midjourney/handle_midjourney.rs | 151 +---- .../commands/download/download_url_command.rs | 18 +- .../download/get_download_path_command.rs | 46 ++ .../src/core/commands/download/mod.rs | 1 + .../core/commands/enqueue/generate_error.rs | 3 + ...router_image_request_to_artcraft_prompt.rs | 1 + ...router_video_request_to_artcraft_prompt.rs | 1 + .../generate_image/generate_image_command.rs | 11 +- .../providers/artcraft_router/mod.rs | 2 +- .../providers/midjourney/handle_midjourney.rs | 182 ++---- .../providers/midjourney/mod.rs | 2 +- .../generate/generate_image/providers/mod.rs | 3 +- .../tauri_generate_image_request.rs | 1 + .../core/commands/generate/omni/request.rs | 28 +- .../state/app_preferences/app_preferences.rs | 8 + .../app_preferences_serializable.rs | 56 +- .../src/core/state/app_preferences/mod.rs | 3 +- .../preferred_download_filename.rs | 351 ++++++++++ .../state/data_dir/app_credentials_dir.rs | 4 + .../handlers/fal/handle_fal_complete.rs | 7 +- .../artcraft/src/core/utils/auto_download.rs | 309 +++++++++ .../download_url_to_user_download_dir.rs | 117 ++-- crates/desktop/artcraft/src/core/utils/mod.rs | 3 +- crates/desktop/artcraft/src/lib.rs | 2 + .../grok_image_websocket_thread.rs | 7 +- .../grok_video_task_polling_thread.rs | 5 +- .../midjourney_clear_credentials_command.rs | 35 +- .../midjourney/state/midjourney_credential.rs | 84 +++ .../state/midjourney_credential_manager.rs | 466 +++++++++----- .../src/services/midjourney/state/mod.rs | 3 +- .../threads/midjourney_long_polling_thread.rs | 602 +++++++----------- .../src/services/midjourney/threads/mod.rs | 3 +- .../threads/upload_midjourney_batch.rs | 202 ++++++ ...extract_midjourney_user_id_from_cookies.rs | 62 ++ .../utils/midjourney_browser_profile.rs | 21 + .../src/services/midjourney/utils/mod.rs | 4 +- .../extract_midjourney_webview_cookies.rs | 49 +- .../windows/midjourney_login_window_thread.rs | 178 +----- .../windows/open_midjourney_login_window.rs | 51 +- .../helpers/handle_successful_generations.rs | 6 + .../handle_storyteller_successful_job.rs | 9 + .../storyteller_task_polling_thread.rs | 5 +- .../storyteller_billing_window_thread.rs | 1 - .../threads/worldlabs_marble_task_polling.rs | 4 + crates/lib/browser_emulation/Cargo.toml | 17 + .../browser_emulation/src/browser_profile.rs | 225 +++++++ crates/lib/browser_emulation/src/lib.rs | 36 ++ crates/lib/cookie_store_wrapper/Cargo.toml | 37 ++ .../cookie_store_wrapper/src/change_log.rs | 80 +++ .../cookie_store_wrapper/src/cookie_store.rs | 329 ++++++++++ crates/lib/cookie_store_wrapper/src/lib.rs | 36 ++ .../cookie_store_wrapper/src/serde_support.rs | 312 +++++++++ .../src/shared_cookie_store.rs | 78 +++ .../src/wreq_cookie_provider.rs | 60 ++ .../DesktopGenerationGallery.tsx | 2 +- .../generation-feed/desktopMediaActions.ts | 42 +- frontend/libs/api/src/lib/LocalApi.ts | 36 +- .../gallery-modal/src/lib/gallery-modal.tsx | 2 +- .../lightbox-modal/src/lib/lightbox-modal.tsx | 4 +- .../src/lib/classy-model-selector.spec.tsx | 104 +++ .../components/popover/src/lib/popover.tsx | 19 +- .../promptbox/src/lib/PromptBoxImage.tsx | 15 +- .../lib/panes/DownloadsSettingsPane.spec.tsx | 106 +++ .../src/lib/panes/DownloadsSettingsPane.tsx | 254 ++++++++ .../src/lib/panes/MiscSettingsPane.tsx | 127 ---- .../settings-modal/src/lib/settings-modal.tsx | 7 +- .../libs/components/switch/src/lib/switch.tsx | 9 +- frontend/libs/model-list/src/index.ts | 2 + .../lib/classes/imageModelForProvider.spec.ts | 24 + .../src/lib/classes/imageModelForProvider.ts | 27 + .../model-list/src/lib/lists/ImageModels.ts | 2 +- .../lib/loader/buildModelsFromListing.spec.ts | 11 + .../src/lib/loader/buildModelsFromListing.ts | 5 +- .../tauri-api/src/lib/download/DownloadUrl.ts | 14 +- .../src/lib/generate/GenerateImage.ts | 1 + .../src/lib/generate/GenerateOmni.spec.ts | 6 + .../src/lib/settings/GetAppPreferences.ts | 8 + .../src/lib/settings/UpdateAppPreference.ts | 6 +- 126 files changed, 5503 insertions(+), 1566 deletions(-) create mode 100644 crates/api_clients/artcraft/artcraft_router/src/client/router_midjourney_client.rs create mode 100644 crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/build.rs create mode 100644 crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/cost.rs create mode 100644 crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/mod.rs create mode 100644 crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/request.rs delete mode 100644 crates/api_clients/midjourney_client/src/client/midjourney_websocket_client.rs create mode 100644 crates/api_clients/midjourney_client/src/client/websocket/job_progress.rs create mode 100644 crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs create mode 100644 crates/api_clients/midjourney_client/src/client/websocket/midjourney_ws_event.rs create mode 100644 crates/api_clients/midjourney_client/src/client/websocket/mod.rs create mode 100644 crates/api_clients/midjourney_client/src/client/websocket/open_midjourney_websocket.rs create mode 100644 crates/api_clients/midjourney_client/src/endpoints/prompt_session_log.rs create mode 100644 crates/api_clients/midjourney_client/src/endpoints/refresh_token.rs create mode 100644 crates/desktop/artcraft/src/core/commands/download/get_download_path_command.rs create mode 100644 crates/desktop/artcraft/src/core/state/app_preferences/preferred_download_filename.rs create mode 100644 crates/desktop/artcraft/src/core/utils/auto_download.rs create mode 100644 crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential.rs create mode 100644 crates/desktop/artcraft/src/services/midjourney/threads/upload_midjourney_batch.rs create mode 100644 crates/desktop/artcraft/src/services/midjourney/utils/extract_midjourney_user_id_from_cookies.rs create mode 100644 crates/desktop/artcraft/src/services/midjourney/utils/midjourney_browser_profile.rs create mode 100644 crates/lib/browser_emulation/Cargo.toml create mode 100644 crates/lib/browser_emulation/src/browser_profile.rs create mode 100644 crates/lib/browser_emulation/src/lib.rs create mode 100644 crates/lib/cookie_store_wrapper/Cargo.toml create mode 100644 crates/lib/cookie_store_wrapper/src/change_log.rs create mode 100644 crates/lib/cookie_store_wrapper/src/cookie_store.rs create mode 100644 crates/lib/cookie_store_wrapper/src/lib.rs create mode 100644 crates/lib/cookie_store_wrapper/src/serde_support.rs create mode 100644 crates/lib/cookie_store_wrapper/src/shared_cookie_store.rs create mode 100644 crates/lib/cookie_store_wrapper/src/wreq_cookie_provider.rs create mode 100644 frontend/libs/components/model-selector/src/lib/classy-model-selector.spec.tsx create mode 100644 frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.spec.tsx create mode 100644 frontend/libs/components/settings-modal/src/lib/panes/DownloadsSettingsPane.tsx create mode 100644 frontend/libs/model-list/src/lib/classes/imageModelForProvider.spec.ts create mode 100644 frontend/libs/model-list/src/lib/classes/imageModelForProvider.ts diff --git a/Cargo.lock b/Cargo.lock index 215ce04fc0..50b344b749 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -188,10 +188,12 @@ dependencies = [ "artcraft_client", "artcraft_router", "base64 0.22.1", + "browser_emulation", "build_metadata", "bytes", "chrono", "cookie_store 0.0.1", + "cookie_store_wrapper", "directories", "dirs", "discord-rich-presence", @@ -231,6 +233,7 @@ dependencies = [ "tokens", "tokio", "tokio-util", + "toml 1.0.7+spec-1.1.0", "url", "url_utils", "uuid_utils", @@ -286,12 +289,14 @@ dependencies = [ "anyhow", "artcraft_api_defs", "artcraft_client", + "browser_emulation", "enums", "fal_client", "gmicloud_client", "grok_api_client", "kinovi_web_client", "log", + "midjourney_client", "reqwest 0.12.23", "serde", "serde_derive", @@ -762,6 +767,14 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "browser_emulation" +version = "0.0.1" +dependencies = [ + "wreq", + "wreq-util", +] + [[package]] name = "bstr" version = "1.12.1" @@ -1058,6 +1071,33 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -1172,6 +1212,22 @@ dependencies = [ "url", ] +[[package]] +name = "cookie_store_wrapper" +version = "0.0.2" +dependencies = [ + "chrono", + "cookie_store 0.21.1", + "log", + "serde", + "serde_derive", + "serde_json", + "time", + "toml 1.0.7+spec-1.1.0", + "url", + "wreq", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -3839,8 +3895,10 @@ name = "midjourney_client" version = "0.0.1" dependencies = [ "anyhow", + "browser_emulation", + "ciborium", "cloudflare_errors", - "cookie_store 0.0.1", + "cookie_store_wrapper", "dom_query 0.20.2", "errors", "filesys", @@ -4099,7 +4157,7 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "681030a937600a36906c185595136d26abfebb4aa9c65701cefcaf8578bb982b" dependencies = [ - "proc-macro-crate 2.0.0", + "proc-macro-crate 3.5.0", "proc-macro2", "quote", "syn 2.0.117", diff --git a/Cargo.toml b/Cargo.toml index 12fed20259..ee09b17003 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,11 +16,13 @@ members = [ "crates/api_clients/worldlabs_api_client", "crates/api_clients/worldlabs_consumer_client", "crates/desktop/artcraft", + "crates/lib/browser_emulation", "crates/lib/build_metadata", "crates/lib/caching/memory_store", "crates/lib/cloudflare_errors", "crates/lib/cloudflare_mitigation", "crates/lib/cookie_store", + "crates/lib/cookie_store_wrapper", "crates/lib/data/primitives", "crates/lib/data/web_base64", "crates/lib/errors", @@ -48,6 +50,8 @@ publish = false artcraft_api_defs = { path = "crates/api_clients/artcraft/artcraft_api_defs" } artcraft_client = { path = "crates/api_clients/artcraft/artcraft_client" } artcraft_router = { path = "crates/api_clients/artcraft/artcraft_router" } +browser_emulation = { path = "crates/lib/browser_emulation" } +cookie_store_wrapper = { path = "crates/lib/cookie_store_wrapper" } build_metadata = { path = "crates/lib/build_metadata" } cloudflare_errors = { path = "crates/lib/cloudflare_errors" } cloudflare_mitigation = { path = "crates/lib/cloudflare_mitigation" } 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 index 193b9575f7..e4902effcf 100644 --- 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 @@ -48,6 +48,15 @@ pub async fn upload_image_media_file_from_file>( args: UploadImageFromFileArgs<'_, P>, ) -> Result { + upload_image_media_file_from_file_with_idempotency_token(args, &generate_random_uuid()).await +} + +/// Upload using a durable idempotency token so a retry after a lost response +/// cannot create a second media file. +pub async fn upload_image_media_file_from_file_with_idempotency_token>( + args: UploadImageFromFileArgs<'_, P>, + idempotency_token: &str, +) -> Result { validate_args(&args)?; let url = get_route(args.api_host); @@ -66,7 +75,7 @@ pub async fn upload_image_media_file_from_file>( let file_name = args.path.as_ref().file_name() .and_then(|n| n.to_str()).unwrap_or("file").to_string(); let mut form = Form::new() - .text("uuid_idempotency_token", generate_random_uuid()) + .text("uuid_idempotency_token", idempotency_token.to_string()) .part("file", Part::bytes(file_bytes).file_name(file_name)); if args.is_intermediate_system_file { diff --git a/crates/api_clients/artcraft/artcraft_router/Cargo.toml b/crates/api_clients/artcraft/artcraft_router/Cargo.toml index cb0b82c726..3fd929faee 100644 --- a/crates/api_clients/artcraft/artcraft_router/Cargo.toml +++ b/crates/api_clients/artcraft/artcraft_router/Cargo.toml @@ -9,6 +9,8 @@ name = "artcraft_router" path = "src/lib.rs" [dependencies] +midjourney_client.workspace = true +browser_emulation.workspace = true # Internal artcraft_api_defs.workspace = true diff --git a/crates/api_clients/artcraft/artcraft_router/src/api/router_provider.rs b/crates/api_clients/artcraft/artcraft_router/src/api/router_provider.rs index 69f736d47f..2b5a1900cf 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/api/router_provider.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/api/router_provider.rs @@ -5,6 +5,7 @@ use serde_derive::{Deserialize, Serialize}; #[serde(rename_all = "snake_case")] pub enum RouterProvider { Artcraft, + Midjourney, Fal, GmiCloud, GrokApi, diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs b/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs index 3d317551fb..4f2816ef1c 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/mod.rs @@ -9,3 +9,5 @@ pub mod router_gmicloud_client; pub mod router_grok_api_client; pub mod router_kinovi_web_client; pub mod router_worldlabs_client; + +pub mod router_midjourney_client; diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client.rs b/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client.rs index 49e8c5a44d..8dc60cd0a2 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client.rs @@ -1,3 +1,4 @@ +use crate::client::router_midjourney_client::RouterMidjourneyClient; use crate::client::router_artcraft_client::RouterArtcraftClient; use crate::client::router_fal_client::RouterFalClient; use crate::client::router_gmicloud_client::RouterGmiCloudClient; @@ -7,6 +8,7 @@ use crate::client::router_worldlabs_client::RouterWorldLabsClient; use crate::errors::client_error::{ClientError, ClientType}; pub struct MultiRouterClient { + pub(crate) midjourney_client: Option, pub(crate) artcraft_client: Option, pub(crate) fal_client: Option, pub(crate) gmicloud_client: Option, @@ -16,6 +18,11 @@ pub struct MultiRouterClient { } impl MultiRouterClient { + pub fn get_midjourney_client_ref(&self) -> Result<&RouterMidjourneyClient, ClientError> { + self.midjourney_client.as_ref() + .ok_or(ClientError::ClientNotConfigured(ClientType::Midjourney)) + } + pub fn get_artcraft_client_ref(&self) -> Result<&RouterArtcraftClient, ClientError> { self.artcraft_client.as_ref() .ok_or(ClientError::ClientNotConfigured(ClientType::Artcraft)) diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client_builder.rs b/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client_builder.rs index 2a7e947e42..ea660a5b87 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client_builder.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/multi_router_client_builder.rs @@ -1,3 +1,4 @@ +use crate::client::router_midjourney_client::RouterMidjourneyClient; use crate::client::multi_router_client::MultiRouterClient; use crate::client::router_artcraft_client::RouterArtcraftClient; use crate::client::router_fal_client::RouterFalClient; @@ -7,6 +8,7 @@ use crate::client::router_kinovi_web_client::RouterKinoviWebClient; use crate::client::router_worldlabs_client::RouterWorldLabsClient; pub struct MultiRouterClientBuilder { + midjourney_client: Option, artcraft_client: Option, fal_client: Option, gmicloud_client: Option, @@ -19,6 +21,7 @@ impl MultiRouterClientBuilder { pub fn new() -> Self { Self { artcraft_client: None, + midjourney_client: None, fal_client: None, gmicloud_client: None, grok_api_client: None, @@ -57,9 +60,15 @@ impl MultiRouterClientBuilder { self } + pub fn set_midjourney_client(mut self, client: RouterMidjourneyClient) -> Self { + self.midjourney_client = Some(client); + self + } + pub fn build(self) -> MultiRouterClient { MultiRouterClient { artcraft_client: self.artcraft_client, + midjourney_client: self.midjourney_client, fal_client: self.fal_client, gmicloud_client: self.gmicloud_client, grok_api_client: self.grok_api_client, diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/router_client.rs b/crates/api_clients/artcraft/artcraft_router/src/client/router_client.rs index 16b48568a4..1774b96b8f 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/client/router_client.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/client/router_client.rs @@ -1,3 +1,4 @@ +use crate::client::router_midjourney_client::RouterMidjourneyClient; use crate::client::multi_router_client::MultiRouterClient; use crate::client::router_artcraft_client::RouterArtcraftClient; use crate::client::router_fal_client::RouterFalClient; @@ -10,6 +11,7 @@ use crate::errors::client_error::{ClientError, ClientType}; pub enum RouterClient { Multi(MultiRouterClient), Artcraft(RouterArtcraftClient), + Midjourney(RouterMidjourneyClient), Fal(RouterFalClient), GmiCloud(RouterGmiCloudClient), GrokApi(RouterGrokApiClient), @@ -18,6 +20,14 @@ pub enum RouterClient { } impl RouterClient { + pub fn get_midjourney_client_ref(&self) -> Result<&RouterMidjourneyClient, ClientError> { + match self { + Self::Midjourney(client) => Ok(client), + Self::Multi(multi) => multi.get_midjourney_client_ref(), + _ => Err(ClientError::ClientNotConfigured(ClientType::Midjourney)), + } + } + pub fn get_artcraft_client_ref(&self) -> Result<&RouterArtcraftClient, ClientError> { match self { RouterClient::Artcraft(client) => Ok(client), diff --git a/crates/api_clients/artcraft/artcraft_router/src/client/router_midjourney_client.rs b/crates/api_clients/artcraft/artcraft_router/src/client/router_midjourney_client.rs new file mode 100644 index 0000000000..a1588df3f2 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_router/src/client/router_midjourney_client.rs @@ -0,0 +1,24 @@ +use browser_emulation::browser_profile::BrowserProfile; +use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId; + +/// The first-party (cookie-session) Midjourney client used by the router. +/// +/// Unlike the Kinovi/Seedance2Pro Midjourney path (which is backend-billed), +/// this drives the user's own logged-in midjourney.com session directly via +/// captured cookies. The `browser` profile MUST match the browser that +/// captured the cookies — Cloudflare's `cf_clearance` cookie is bound to the +/// User-Agent. +pub struct RouterMidjourneyClient { + pub(crate) cookie_header: String, + + /// Needed to form the `singleplayer_{user_id}` channel id on submit. + pub(crate) user_id: MidjourneyUserId, + + pub(crate) browser: BrowserProfile, +} + +impl RouterMidjourneyClient { + pub fn new(cookie_header: String, user_id: MidjourneyUserId, browser: BrowserProfile) -> Self { + Self { cookie_header, user_id, browser } + } +} diff --git a/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs b/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs index 88507750ec..61402662c9 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs @@ -64,6 +64,7 @@ impl From for ArtcraftRouterError { impl From for ArtcraftRouterError { fn from(error: ProviderError) -> Self { let is_billing_error = match &error { + ProviderError::MidjourneySubscriptionRequired(_) => true, ProviderError::Fal(FalErrorPlus::FalBillingError(_)) => true, ProviderError::KinoviWeb(KinoviWebError::ApiSpecific(KinoviWebSpecificApiError::BillingError { .. })) => true, ProviderError::Storyteller(StorytellerError::Api(ApiError::PaymentRequired(_))) => true, diff --git a/crates/api_clients/artcraft/artcraft_router/src/errors/client_error.rs b/crates/api_clients/artcraft/artcraft_router/src/errors/client_error.rs index 94733d4b55..9f81c44500 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/errors/client_error.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/errors/client_error.rs @@ -6,6 +6,7 @@ use tokens::tokens::media_files::MediaFileToken; #[derive(Debug, Clone, Copy)] pub enum ClientType { Artcraft, + Midjourney, Fal, GmiCloud, GrokApi, @@ -16,6 +17,7 @@ pub enum ClientType { impl Display for ClientType { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { + Self::Midjourney => write!(f, "Midjourney"), Self::Artcraft => write!(f, "Artcraft"), Self::Fal => write!(f, "Fal"), Self::GmiCloud => write!(f, "GmiCloud"), diff --git a/crates/api_clients/artcraft/artcraft_router/src/errors/provider_error.rs b/crates/api_clients/artcraft/artcraft_router/src/errors/provider_error.rs index eac4f7038f..1f31849456 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/errors/provider_error.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/errors/provider_error.rs @@ -1,3 +1,4 @@ +use midjourney_client::error::midjourney_error::MidjourneyError; use std::error::Error; use std::fmt::{Display, Formatter}; use artcraft_client::error::storyteller_error::StorytellerError; @@ -9,6 +10,9 @@ use worldlabs_api_client::error::world_labs_error::WorldLabsError; #[derive(Debug)] pub enum ProviderError { + Midjourney(MidjourneyError), + MidjourneySubscriptionRequired(String), + MidjourneySubmitRejected(String), Storyteller(StorytellerError), Fal(FalErrorPlus), GmiCloud(GmiCloudError), @@ -22,6 +26,8 @@ impl Error for ProviderError {} impl Display for ProviderError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { + Self::Midjourney(e) => write!(f, "Midjourney provider error: {}", e), + Self::MidjourneySubscriptionRequired(message) | Self::MidjourneySubmitRejected(message) => write!(f, "{}", message), Self::Storyteller(e) => write!(f, "Storyteller provider error: {}", e), Self::Fal(e) => write!(f, "Fal provider error: {}", e), Self::GmiCloud(e) => write!(f, "GmiCloud provider error: {}", e), diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_request_builder.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_request_builder.rs index 2a92513333..e916d4ec41 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_request_builder.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_request_builder.rs @@ -1,3 +1,4 @@ +use crate::generate::generate_image::providers::midjourney::build::build_midjourney; use crate::api::router_aspect_ratio::RouterAspectRatio; use crate::api::router_image_model::RouterImageModel; use crate::api::router_quality::RouterQuality; @@ -113,6 +114,7 @@ impl GenerateImageRequestBuilder { /// `GenerateVideoRequestBuilder::build2`. pub fn build2(self) -> Result { match (self.provider, self.model) { + (RouterProvider::Midjourney, RouterImageModel::Midjourney7 | RouterImageModel::Midjourney7Niji | RouterImageModel::Midjourney8) => build_midjourney(self), (RouterProvider::Artcraft, RouterImageModel::Flux1Dev) => build_artcraft_flux_1_dev(self), (RouterProvider::Artcraft, RouterImageModel::Flux1Schnell) => build_artcraft_flux_1_schnell(self), (RouterProvider::Artcraft, RouterImageModel::FluxPro11) => build_artcraft_flux_pro_1p1(self), diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs index 643e0ecdbe..e79bfae714 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/generate_image_response.rs @@ -38,6 +38,7 @@ pub struct KinoviWebImageResponsePayload { #[derive(Clone, Debug)] pub enum GenerateImageResponse { + Midjourney(MidjourneyImageResponsePayload), Artcraft(ArtcraftImageResponsePayload), Fal(FalImageResponsePayload), KinoviWeb(KinoviWebImageResponsePayload), @@ -65,3 +66,8 @@ impl GenerateImageResponse { } } } + +#[derive(Clone, Debug)] +pub struct MidjourneyImageResponsePayload { + pub job_id: String, +} diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/image_generation_request.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/image_generation_request.rs index 5eed144b5c..fc9ee7e50c 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/image_generation_request.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/image_generation_request.rs @@ -1,3 +1,5 @@ +use crate::generate::generate_image::providers::midjourney::request::MidjourneyRequestState; +use crate::generate::generate_image::providers::midjourney::cost::MidjourneyCostState; use crate::api::router_provider::RouterProvider; use crate::client::router_client::RouterClient; use crate::errors::artcraft_router_error::ArtcraftRouterError; @@ -92,6 +94,7 @@ use crate::generate::generate_image::providers::kinovi::seedream_5p0_pro::reques #[derive(Clone, Debug)] pub enum ImageGenerationRequest { + MidjourneyRequest(MidjourneyRequestState), // ── Artcraft provider (omni-gen image endpoint) ── ArtcraftFlux1Dev(ArtcraftFlux1DevRequestState), ArtcraftFlux1Schnell(ArtcraftFlux1SchnellRequestState), @@ -145,6 +148,7 @@ pub enum ImageGenerationRequest { impl ImageGenerationRequest { pub fn get_provider(&self) -> RouterProvider { match self { + Self::MidjourneyRequest(_) => RouterProvider::Midjourney, Self::ArtcraftFlux1Dev(_) => RouterProvider::Artcraft, Self::ArtcraftFlux1Schnell(_) => RouterProvider::Artcraft, Self::ArtcraftFluxPro1p1(_) => RouterProvider::Artcraft, @@ -195,6 +199,7 @@ impl ImageGenerationRequest { pub fn estimate_cost(&self) -> Result { match self { + Self::MidjourneyRequest(request) => Ok(MidjourneyCostState::from_request(request).estimate_cost()), // ── Artcraft ── Self::ArtcraftFlux1Dev(request) => { Ok(ArtcraftFlux1DevCostState::from_request(request).estimate_cost()) @@ -291,6 +296,7 @@ impl ImageGenerationRequest { pub async fn send_request(&self, client: &RouterClient) -> Result { match self { + Self::MidjourneyRequest(request) => request.send(client.get_midjourney_client_ref()?).await, // ── Artcraft (omni-gen image endpoint) ── Self::ArtcraftFlux1Dev(request) => { let artcraft_client = client.get_artcraft_client_ref()?; diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/build.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/build.rs new file mode 100644 index 0000000000..d0c7ab8d26 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/build.rs @@ -0,0 +1,130 @@ +use crate::api::router_image_model::RouterImageModel; +use crate::api::image_list_ref::ImageListRef; +use crate::api::router_aspect_ratio::RouterAspectRatio; +use crate::errors::artcraft_router_error::ArtcraftRouterError; +use crate::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; +use crate::generate::generate_image::image_generation_draft_or_request::ImageGenerationDraftOrRequest; +use crate::generate::generate_image::image_generation_request::ImageGenerationRequest; +use crate::generate::generate_image::providers::midjourney::request::MidjourneyRequestState; + +/// The Midjourney model version flag appended to every prompt. +const MIDJOURNEY_VERSION_FLAG: &str = "--v 8.2"; + +pub fn build_midjourney(builder: GenerateImageRequestBuilder) -> Result { + // Image references are not yet supported on the first-party cookie path. + if has_image_inputs(builder.image_inputs.as_ref()) { + return Err(ArtcraftRouterError::UnsupportedProviderAndModelForNewApi("First-party Midjourney does not yet support image references".to_string())); + } + + if builder.image_batch_count.is_some_and(|count| count != 4) { + return Err(ArtcraftRouterError::InvalidInput("Direct Midjourney generates four images per job".to_string())); + } + let version = match builder.model { + RouterImageModel::Midjourney7 => "--v 7", + RouterImageModel::Midjourney7Niji => "--niji 7", + RouterImageModel::Midjourney8 => MIDJOURNEY_VERSION_FLAG, + _ => return Err(ArtcraftRouterError::UnsupportedModel(format!("{:?}", builder.model))), + }; + let base_prompt = builder.prompt.unwrap_or_default(); + if base_prompt.trim().is_empty() { + return Err(ArtcraftRouterError::InvalidInput("A Midjourney prompt is required".to_string())); + } + let prompt = compose_prompt(&base_prompt, builder.aspect_ratio, version); + + Ok(ImageGenerationDraftOrRequest::Request(ImageGenerationRequest::MidjourneyRequest(MidjourneyRequestState { prompt }))) +} + +/// Compose the final Midjourney prompt string, appending the aspect-ratio flag +/// (when the ratio maps to one Midjourney supports) and the version flag. +fn compose_prompt(base_prompt: &str, aspect_ratio: Option, version: &str) -> String { + let mut prompt = base_prompt.trim().to_string(); + + if let Some(ratio) = aspect_ratio_flag(aspect_ratio) { + prompt.push_str(" --ar "); + prompt.push_str(ratio); + } + + prompt.push(' '); + prompt.push_str(version); + prompt +} + +/// Maps a router aspect ratio to Midjourney's `--ar W:H` value. Returns `None` +/// for `Auto`/unset (Midjourney defaults to 1:1). +fn aspect_ratio_flag(aspect_ratio: Option) -> Option<&'static str> { + match aspect_ratio { + None | Some(RouterAspectRatio::Auto) | Some(RouterAspectRatio::Auto2k) | Some(RouterAspectRatio::Auto3k) | Some(RouterAspectRatio::Auto4k) => None, + + Some(RouterAspectRatio::Square) | Some(RouterAspectRatio::SquareHd) => Some("1:1"), + Some(RouterAspectRatio::WideThreeByTwo) => Some("3:2"), + Some(RouterAspectRatio::WideFourByThree) => Some("4:3"), + Some(RouterAspectRatio::WideFiveByFour) => Some("5:4"), + Some(RouterAspectRatio::WideSixteenByNine) | Some(RouterAspectRatio::Wide) => Some("16:9"), + Some(RouterAspectRatio::WideTwentyOneByNine) => Some("21:9"), + Some(RouterAspectRatio::TallTwoByThree) => Some("2:3"), + Some(RouterAspectRatio::TallThreeByFour) => Some("3:4"), + Some(RouterAspectRatio::TallFourByFive) => Some("4:5"), + Some(RouterAspectRatio::TallNineBySixteen) | Some(RouterAspectRatio::Tall) => Some("9:16"), + Some(RouterAspectRatio::TallNineByTwentyOne) => Some("9:21"), + } +} + +fn has_image_inputs(image_inputs: Option<&ImageListRef>) -> bool { + image_inputs.is_some_and(|refs| match refs { + ImageListRef::Urls(urls) => !urls.is_empty(), + ImageListRef::MediaFileTokens(tokens) => !tokens.is_empty(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::router_image_model::RouterImageModel; + use crate::api::router_provider::RouterProvider; + use crate::client::generation_mode_mismatch_strategy::GenerationModeMismatchStrategy; + use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy; + + #[test] + fn appends_version_flag_and_no_ar_by_default() { + let prompt = built_prompt(base_builder()); + assert_eq!(prompt, "pirate ship in the city --v 8.2"); + } + + #[test] + fn appends_aspect_ratio_flag() { + let builder = GenerateImageRequestBuilder { aspect_ratio: Some(RouterAspectRatio::WideSixteenByNine), ..base_builder() }; + assert_eq!(built_prompt(builder), "pirate ship in the city --ar 16:9 --v 8.2"); + } + + #[test] + fn image_inputs_are_rejected_for_now() { + let builder = GenerateImageRequestBuilder { image_inputs: Some(ImageListRef::Urls(vec!["https://example.com/ref.png".to_string()])), ..base_builder() }; + assert!(build_midjourney(builder).is_err()); + } + #[test] + fn chooses_the_requested_midjourney_version() { + for (model, flag) in [(RouterImageModel::Midjourney7, "--v 7"), (RouterImageModel::Midjourney7Niji, "--niji 7"), (RouterImageModel::Midjourney8, "--v 8.2")] { + let prompt = built_prompt(GenerateImageRequestBuilder { model, ..base_builder() }); + assert!(prompt.ends_with(flag)); + } + } + + #[test] + fn rejects_invalid_batch_counts_and_empty_prompts() { + for count in [0, 1, 2, 5] { + assert!(build_midjourney(GenerateImageRequestBuilder { image_batch_count: Some(count), ..base_builder() }).is_err()); + } + assert!(build_midjourney(GenerateImageRequestBuilder { prompt: Some(" ".into()), ..base_builder() }).is_err()); + } + + fn base_builder() -> GenerateImageRequestBuilder { + GenerateImageRequestBuilder { model: RouterImageModel::Midjourney8, provider: RouterProvider::Midjourney, prompt: Some("pirate ship in the city".to_string()), image_inputs: None, resolution: None, aspect_ratio: None, quality: None, image_batch_count: None, horizontal_angle: None, vertical_angle: None, zoom: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, generation_mode_mismatch_strategy: Some(GenerationModeMismatchStrategy::GenerateAnyway), idempotency_token: None } + } + + fn built_prompt(builder: GenerateImageRequestBuilder) -> String { + match build_midjourney(builder).expect("build") { + ImageGenerationDraftOrRequest::Request(ImageGenerationRequest::MidjourneyRequest(r)) => r.prompt, + _ => panic!("expected a direct Midjourney request"), + } + } +} diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/cost.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/cost.rs new file mode 100644 index 0000000000..2bb234b080 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/cost.rs @@ -0,0 +1,15 @@ +use crate::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate; +use crate::generate::generate_image::providers::midjourney::request::MidjourneyRequestState; + +/// Direct Midjourney uses the user's paid subscription; it consumes no ArtCraft credits. +pub struct MidjourneyCostState; + +impl MidjourneyCostState { + pub fn from_request(_request: &MidjourneyRequestState) -> Self { + Self + } + + pub fn estimate_cost(&self) -> ImageGenerationCostEstimate { + ImageGenerationCostEstimate { cost_in_credits: Some(0), cost_in_usd_cents: None, is_free: false, is_unlimited: false, is_rate_limited: true, has_watermark: false, failures_are_refunded: None } + } +} diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/mod.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/mod.rs new file mode 100644 index 0000000000..b2208da473 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/mod.rs @@ -0,0 +1,3 @@ +pub mod build; +pub mod request; +pub mod cost; diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/request.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/request.rs new file mode 100644 index 0000000000..cfdd7c7a7b --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/midjourney/request.rs @@ -0,0 +1,41 @@ +use midjourney_client::recipes::channel_id::ChannelId; +use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageArgs, TextToImageRequest, TextToImageResponse}; + +use crate::client::router_midjourney_client::RouterMidjourneyClient; +use crate::errors::artcraft_router_error::ArtcraftRouterError; +use crate::errors::provider_error::ProviderError; +use crate::generate::generate_image::generate_image_response::{GenerateImageResponse, MidjourneyImageResponsePayload}; + +/// A ready-to-send direct Midjourney request. The `prompt` is already +/// composed with the selected aspect ratio and model version. +#[derive(Debug, Clone)] +pub struct MidjourneyRequestState { + pub prompt: String, +} + +impl MidjourneyRequestState { + pub async fn send(&self, client: &RouterMidjourneyClient) -> Result { + let channel_id = ChannelId::UserId(client.user_id.clone()); + + let response = text_to_image(TextToImageArgs { request: TextToImageRequest { prompt: &self.prompt, channel_id: &channel_id }, cookie_header: &client.cookie_header, hostname: None, browser: Some(client.browser.clone()) }).await.map_err(|err| ArtcraftRouterError::Provider(ProviderError::Midjourney(err)))?; + + if let Some(job_id) = response.maybe_job_id.clone() { + return Ok(GenerateImageResponse::Midjourney(MidjourneyImageResponsePayload { job_id })); + } + + // No job id — classify the rejection. Surface Midjourney's own + // human-readable message(s), not a debug dump. + let detail = midjourney_error_detail(&response); + if response.is_subscription_required() { + return Err(ArtcraftRouterError::Provider(ProviderError::MidjourneySubscriptionRequired(detail))); + } + Err(ArtcraftRouterError::Provider(ProviderError::MidjourneySubmitRejected(detail))) + } +} + +/// Join Midjourney's failure messages into a single human-readable string. +fn midjourney_error_detail(response: &TextToImageResponse) -> String { + let detail = response.errors().iter().filter_map(|error| error.message.clone()).collect::>().join("; "); + + if detail.is_empty() { "Midjourney rejected the request with no detail".to_string() } else { detail } +} diff --git a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/mod.rs b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/mod.rs index 1adb32b1fa..ed201ade57 100644 --- a/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/mod.rs +++ b/crates/api_clients/artcraft/artcraft_router/src/generate/generate_image/providers/mod.rs @@ -1,3 +1,5 @@ pub mod artcraft; pub mod fal; pub mod kinovi; + +pub mod midjourney; diff --git a/crates/api_clients/midjourney_client/Cargo.toml b/crates/api_clients/midjourney_client/Cargo.toml index 177b4f403c..d0392b7de2 100644 --- a/crates/api_clients/midjourney_client/Cargo.toml +++ b/crates/api_clients/midjourney_client/Cargo.toml @@ -9,10 +9,10 @@ name = "midjourney_client" path = "src/lib.rs" [dependencies] - # Internal +browser_emulation.workspace = true cloudflare_errors.workspace = true -cookie_store.workspace = true +cookie_store_wrapper = { workspace = true, features = ["wreq"] } errors.workspace = true # External @@ -29,6 +29,9 @@ wreq.workspace = true # NB: We need to parse HTML to get out the websocket token dom_query = "0.20.2" +# NB: Midjourney's websocket sends server->client frames as CBOR. +ciborium = "0.2" + [dev-dependencies] filesys.workspace = true test_data.workspace = true diff --git a/crates/api_clients/midjourney_client/src/client/midjourney_websocket_client.rs b/crates/api_clients/midjourney_client/src/client/midjourney_websocket_client.rs deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/crates/api_clients/midjourney_client/src/client/mod.rs b/crates/api_clients/midjourney_client/src/client/mod.rs index 233febf33d..5faec3d6e1 100644 --- a/crates/api_clients/midjourney_client/src/client/mod.rs +++ b/crates/api_clients/midjourney_client/src/client/mod.rs @@ -1,2 +1,2 @@ pub mod midjourney_hostname; -pub mod midjourney_websocket_client; +pub mod websocket; diff --git a/crates/api_clients/midjourney_client/src/client/websocket/job_progress.rs b/crates/api_clients/midjourney_client/src/client/websocket/job_progress.rs new file mode 100644 index 0000000000..aad48f0c52 --- /dev/null +++ b/crates/api_clients/midjourney_client/src/client/websocket/job_progress.rs @@ -0,0 +1,117 @@ +/// A single in-progress preview image carried on the websocket during a job. +/// +/// Midjourney streams four images (indices 0-3) that refine over successive +/// diffusion steps. Each websocket progress frame may carry one preview per +/// image for the current step; the previews grow sharper as the step count +/// climbs. The final, full-resolution image is NOT delivered here — it is +/// fetched over HTTP once the job reaches `completed` (see the `imagine` +/// endpoint). +#[derive(Clone)] +pub struct JobStepImage { + /// Which of the four grid images this is, in `[0, 4)`. + pub image_index: u8, + + /// The diffusion step this preview was rendered at (larger = more refined). + pub step: u32, + + /// The image encoding. Observed to be JPEG in practice. + pub format: MidjourneyImageFormat, + + /// The raw encoded image bytes. + pub bytes: Vec, +} + +/// The encoding of a websocket preview image, taken from the file extension in +/// the frame and cross-checked against the leading magic bytes. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum MidjourneyImageFormat { + Jpeg, + Webp, + Png, + Other(String), +} + +impl JobStepImage { + /// Parse one `files` entry, whose key looks like `"0_step_5.jpeg"`. + /// Returns `None` if the key does not match the expected shape. + pub fn from_file_entry(key: &str, bytes: Vec) -> Option { + // "0_step_5.jpeg" -> index "0", rest "step_5.jpeg" + let (index_str, rest) = key.split_once('_')?; + let image_index = index_str.parse::().ok()?; + + // "step_5.jpeg" -> "5.jpeg" + let rest = rest.strip_prefix("step_")?; + + // "5.jpeg" -> step "5", extension "jpeg" + let (step_str, extension) = rest.split_once('.')?; + let step = step_str.parse::().ok()?; + + let format = MidjourneyImageFormat::infer(extension, &bytes); + + Some(Self { image_index, step, format, bytes }) + } +} + +impl MidjourneyImageFormat { + /// Prefer the magic bytes (authoritative), falling back to the extension. + fn infer(extension: &str, bytes: &[u8]) -> Self { + if bytes.len() >= 4 && bytes[0..3] == [0xFF, 0xD8, 0xFF] { + return Self::Jpeg; + } + if bytes.len() >= 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" { + return Self::Webp; + } + if bytes.len() >= 8 && bytes[0..8] == [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A] { + return Self::Png; + } + match extension.to_ascii_lowercase().as_str() { + "jpeg" | "jpg" => Self::Jpeg, + "webp" => Self::Webp, + "png" => Self::Png, + other => Self::Other(other.to_string()), + } + } + + pub fn extension(&self) -> &str { + match self { + Self::Jpeg => "jpeg", + Self::Webp => "webp", + Self::Png => "png", + Self::Other(ext) => ext, + } + } +} + +impl std::fmt::Debug for JobStepImage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("JobStepImage").field("image_index", &self.image_index).field("step", &self.step).field("format", &self.format).field("bytes", &format!("<{} bytes>", self.bytes.len())).finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const JPEG_MAGIC: &[u8] = &[0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10]; + + #[test] + fn parses_step_file_key() { + let image = JobStepImage::from_file_entry("2_step_14.jpeg", JPEG_MAGIC.to_vec()).unwrap(); + assert_eq!(image.image_index, 2); + assert_eq!(image.step, 14); + assert_eq!(image.format, MidjourneyImageFormat::Jpeg); + } + + #[test] + fn magic_bytes_override_extension() { + let image = JobStepImage::from_file_entry("0_step_0.png", JPEG_MAGIC.to_vec()).unwrap(); + assert_eq!(image.format, MidjourneyImageFormat::Jpeg); + } + + #[test] + fn rejects_malformed_keys() { + assert!(JobStepImage::from_file_entry("garbage", vec![]).is_none()); + assert!(JobStepImage::from_file_entry("x_step_1.jpeg", vec![]).is_none()); + assert!(JobStepImage::from_file_entry("0_notstep_1.jpeg", vec![]).is_none()); + } +} 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 new file mode 100644 index 0000000000..a0f2b553ab --- /dev/null +++ b/crates/api_clients/midjourney_client/src/client/websocket/midjourney_websocket.rs @@ -0,0 +1,190 @@ +use crate::client::websocket::midjourney_ws_event::MidjourneyWsEvent; +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; + +/// How many events the broadcast channel buffers. A single job emits ~30 +/// frames; this holds many concurrent jobs before a slow consumer lags. +const EVENT_CHANNEL_CAPACITY: usize = 2048; + +/// A live, thread-safe handle to an open Midjourney websocket. +/// +/// The connection is driven by a background task that owns the raw socket; +/// this handle only talks to it over channels, so it is cheap to `clone` and +/// share across threads. Cloned handles observe the same connection. +/// +/// Reading images: +/// +/// ```ignore +/// let ws = open_midjourney_websocket(/* ... */).await?; +/// // Submit the job over HTTP, then stream its previews: +/// ws.stream_job(&job_id, |event| { +/// if let MidjourneyWsEvent::Progress(progress) = event { +/// for image in &progress.images { +/// // image.bytes is a JPEG preview for image.image_index at image.step +/// } +/// } +/// }).await?; +/// ``` +/// +/// For finer control, call [`Self::events`] to get a raw receiver and +/// [`Self::subscribe_to_job`] / [`Self::send_json`] to drive the protocol +/// yourself. +/// +/// Dropping the last handle ends the background task and closes the socket. +#[derive(Clone)] +pub struct MidjourneyWebSocket { + inner: Arc, +} + +struct WebSocketInner { + user_id: MidjourneyUserId, + outbound: mpsc::UnboundedSender, + events: broadcast::Sender>, + connected: Arc, +} + +impl MidjourneyWebSocket { + /// Take ownership of a freshly upgraded socket, spawn its driver task, and + /// return a handle. The caller is responsible for the `subscribe_to_user` + /// handshake (see `open_midjourney_websocket`). + 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())); + + Self { inner: Arc::new(WebSocketInner { user_id, outbound: outbound_tx, events: events_tx, connected }) } + } + + /// The user this connection is bound to. + pub fn user_id(&self) -> &MidjourneyUserId { + &self.inner.user_id + } + + /// Whether the background task is still running. + pub fn is_connected(&self) -> bool { + self.inner.connected.load(Ordering::SeqCst) + } + + /// A new receiver over the decoded event stream. Each call yields an + /// independent receiver that sees events sent *after* it was created, so + /// subscribe before sending the request whose events you want. + pub fn events(&self) -> broadcast::Receiver> { + self.inner.events.subscribe() + } + + /// Bind this connection to the user (the opening handshake). Normally done + /// for you by `open_midjourney_websocket`. + pub fn subscribe_to_user(&self) -> Result<(), MidjourneyClientError> { + self.send_json(&serde_json::json!({ "type": "subscribe_to_user" })) + } + + /// Ask the server to stream progress for a job (after it was submitted over + /// HTTP). Progress frames, including preview images, follow. + pub fn subscribe_to_job(&self, job_id: &str) -> Result<(), MidjourneyClientError> { + self.send_json(&serde_json::json!({ "type": "subscribe_to_job", "job_id": job_id })) + } + + /// Send an arbitrary JSON command frame. The escape hatch for protocol + /// messages this handle does not model directly (e.g. `room_new_job`). + pub fn send_json(&self, value: &serde_json::Value) -> Result<(), MidjourneyClientError> { + let text = serde_json::to_string(value)?; + self.send_message(Message::text(text)) + } + + /// Subscribe to a job and invoke `on_event` for every event belonging to it, + /// returning once the job completes (or erroring if the socket closes first). + /// + /// This is the ergonomic path for reading images: each `Progress` event + /// carries the preview `images` for the current step, and the call returns + /// when the `Completed` event arrives. + pub async fn stream_job(&self, job_id: &str, mut on_event: F) -> Result<(), MidjourneyClientError> + where + F: FnMut(&MidjourneyWsEvent), + { + let mut receiver = self.events(); + self.subscribe_to_job(job_id)?; + + loop { + match receiver.recv().await { + Ok(event) => { + if event.job_id() != Some(job_id) { + continue; + } + on_event(&event); + if event.is_terminal() { + return Ok(()); + } + }, + Err(broadcast::error::RecvError::Lagged(skipped)) => { + warn!("Midjourney websocket consumer lagged, skipped {} events", skipped); + }, + Err(broadcast::error::RecvError::Closed) => { + return Err(MidjourneyClientError::WebSocketClosed); + }, + } + } + } + + /// Politely close the connection. The background task forwards the close + /// frame and then stops. Dropping the last handle does the same implicitly. + pub fn close(&self) { + let _ = self.send_message(Message::Close(None)); + } + + fn send_message(&self, message: Message) -> Result<(), MidjourneyClientError> { + self.inner.outbound.send(message).map_err(|_| MidjourneyClientError::WebSocketClosed) + } +} + +/// 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) { + loop { + tokio::select! { + maybe_outbound = outbound_rx.recv() => { + match maybe_outbound { + Some(message) => { + if let Err(err) = websocket.send(message).await { + warn!("Midjourney websocket send failed: {}", err); + break; + } + } + // All handles dropped; nothing more will be sent. + None => break, + } + } + maybe_incoming = websocket.recv() => { + match maybe_incoming { + Some(Ok(Message::Binary(bytes))) => { + match MidjourneyWsEvent::from_cbor_frame(&bytes) { + // A send error just means there are no live receivers; that's fine. + Ok(event) => { let _ = events_tx.send(Arc::new(event)); } + Err(err) => warn!("Midjourney websocket frame decode failed: {}", err), + } + } + Some(Ok(Message::Text(text))) => { + debug!("Unexpected text frame from Midjourney websocket: {}", text.as_str()); + } + // Ping/Pong are handled by wreq; ignore anything else non-terminal. + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(_)) => {} + Some(Err(err)) => { + warn!("Midjourney websocket receive error: {}", err); + break; + } + } + } + } + } + + connected.store(false, Ordering::SeqCst); + debug!("Midjourney websocket loop ended."); +} diff --git a/crates/api_clients/midjourney_client/src/client/websocket/midjourney_ws_event.rs b/crates/api_clients/midjourney_client/src/client/websocket/midjourney_ws_event.rs new file mode 100644 index 0000000000..ced2080f04 --- /dev/null +++ b/crates/api_clients/midjourney_client/src/client/websocket/midjourney_ws_event.rs @@ -0,0 +1,231 @@ +use crate::client::websocket::job_progress::JobStepImage; +use ciborium::value::Value as CborValue; + +/// A decoded server-to-client websocket event. +/// +/// Midjourney sends these as CBOR binary frames. The variants cover the +/// handshake reply, job lifecycle notifications, and the streaming progress +/// frames that carry preview images. +#[derive(Clone, Debug)] +pub enum MidjourneyWsEvent { + /// Reply to `subscribe_to_user`; the connection is now bound to this user. + UserSuccess { user_id: String }, + + /// A job was announced in the user's room (echo of `room_new_job`). + RoomNewJob { job_id: Option }, + + /// Reply to `subscribe_to_job`; progress frames for it will now arrive. + JobSuccess { job_id: String }, + + /// A progress update. Carries preview images when the frame includes them. + Progress(JobProgress), + + /// The job finished. The full-resolution result is fetched over HTTP. + Completed { job_id: String, percentage_complete: f64 }, + + /// A frame we recognized as CBOR but do not model. `maybe_kind` is its + /// `type` field when present. + Unknown { maybe_kind: Option }, +} + +/// A single progress frame for a running job. +#[derive(Clone, Debug)] +pub struct JobProgress { + pub job_id: String, + pub maybe_user_id: Option, + + /// e.g. `unqueue`, `start_stage`, `running`. + pub current_status: String, + + pub percentage_complete: f64, + + /// Preview images included in this frame (may be empty). + pub images: Vec, +} + +impl MidjourneyWsEvent { + /// The job this event concerns, if any. `UserSuccess` and unmodeled frames + /// return `None`. + pub fn job_id(&self) -> Option<&str> { + match self { + Self::UserSuccess { .. } | Self::Unknown { .. } => None, + Self::RoomNewJob { job_id } => job_id.as_deref(), + Self::JobSuccess { job_id } => Some(job_id), + Self::Progress(progress) => Some(&progress.job_id), + Self::Completed { job_id, .. } => Some(job_id), + } + } + + /// Whether this event marks the end of a job's progress stream. + pub fn is_terminal(&self) -> bool { + matches!(self, Self::Completed { .. }) + } + + /// Decode one binary websocket frame (CBOR) into an event. + pub fn from_cbor_frame(bytes: &[u8]) -> Result { + let value: CborValue = ciborium::from_reader(bytes).map_err(|err| MidjourneyWsDecodeError(err.to_string()))?; + + let entries = match &value { + CborValue::Map(entries) => entries, + _ => return Err(MidjourneyWsDecodeError("frame was not a CBOR map".to_string())), + }; + + let maybe_kind = get_str(entries, "type"); + + if let Some(kind) = maybe_kind.as_deref() { + match kind { + "user_success" => { + return Ok(Self::UserSuccess { user_id: get_str(entries, "user_id").unwrap_or_default() }); + }, + "room_new_job" => { + return Ok(Self::RoomNewJob { job_id: get_str(entries, "job_id") }); + }, + "job_success" => { + return Ok(Self::JobSuccess { job_id: get_str(entries, "job_id").unwrap_or_default() }); + }, + _ => {}, + } + } + + // Progress frames have no `type`, but do carry `current_status`. + if let Some(current_status) = get_str(entries, "current_status") { + let job_id = get_str(entries, "job_id").unwrap_or_default(); + let percentage_complete = get_f64(entries, "percentage_complete").unwrap_or(0.0); + + if current_status == "completed" { + return Ok(Self::Completed { job_id, percentage_complete }); + } + + return Ok(Self::Progress(JobProgress { job_id, maybe_user_id: get_str(entries, "user_id"), current_status, percentage_complete, images: extract_images(entries) })); + } + + Ok(Self::Unknown { maybe_kind }) + } +} + +/// The error returned when a websocket frame cannot be decoded as CBOR. +#[derive(Debug)] +pub struct MidjourneyWsDecodeError(pub String); + +impl std::fmt::Display for MidjourneyWsDecodeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Failed to decode Midjourney websocket frame: {}", self.0) + } +} + +impl std::error::Error for MidjourneyWsDecodeError {} + +fn extract_images(entries: &[(CborValue, CborValue)]) -> Vec { + let files = match map_get(entries, "files") { + Some(CborValue::Map(files)) => files, + _ => return Vec::new(), + }; + + let mut images = Vec::new(); + for (key, value) in files { + let (CborValue::Text(name), CborValue::Bytes(bytes)) = (key, value) else { + continue; + }; + if let Some(image) = JobStepImage::from_file_entry(name, bytes.clone()) { + images.push(image); + } + } + images.sort_by_key(|image| (image.step, image.image_index)); + images +} + +fn map_get<'a>(entries: &'a [(CborValue, CborValue)], key: &str) -> Option<&'a CborValue> { + entries.iter().find(|(k, _)| matches!(k, CborValue::Text(text) if text == key)).map(|(_, value)| value) +} + +/// Read a string field, accepting both a plain text value and a CBOR +/// tag-37 (RFC 8949 UUID) 16-byte value, which Midjourney uses for some +/// `job_id` / `user_id` fields. +fn get_str(entries: &[(CborValue, CborValue)], key: &str) -> Option { + match map_get(entries, key)? { + CborValue::Text(text) => Some(text.clone()), + CborValue::Tag(37, inner) => match inner.as_ref() { + CborValue::Bytes(bytes) if bytes.len() == 16 => Some(format_uuid(bytes)), + CborValue::Text(text) => Some(text.clone()), + _ => None, + }, + _ => None, + } +} + +fn get_f64(entries: &[(CborValue, CborValue)], key: &str) -> Option { + match map_get(entries, key)? { + CborValue::Float(float) => Some(*float), + CborValue::Integer(integer) => { + let as_i128: i128 = (*integer).into(); + Some(as_i128 as f64) + }, + _ => None, + } +} + +fn format_uuid(bytes: &[u8]) -> String { + let hex: String = bytes.iter().map(|byte| format!("{:02x}", byte)).collect(); + format!("{}-{}-{}-{}-{}", &hex[0..8], &hex[8..12], &hex[12..16], &hex[16..20], &hex[20..32],) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode(value: &CborValue) -> Vec { + let mut buffer = Vec::new(); + ciborium::into_writer(value, &mut buffer).unwrap(); + buffer + } + + fn text(value: &str) -> CborValue { + CborValue::Text(value.to_string()) + } + + #[test] + fn decodes_user_success() { + let frame = encode(&CborValue::Map(vec![(text("type"), text("user_success")), (text("user_id"), text("26c9d38d-1f71-49c7-a356-b29dac58b54c"))])); + + let event = MidjourneyWsEvent::from_cbor_frame(&frame).unwrap(); + let MidjourneyWsEvent::UserSuccess { user_id } = event else { + panic!("expected UserSuccess, got {event:?}"); + }; + assert_eq!(user_id, "26c9d38d-1f71-49c7-a356-b29dac58b54c"); + } + + #[test] + fn decodes_tag37_uuid_job_id() { + let uuid_bytes = vec![0x7d, 0xa8, 0xda, 0x16, 0x69, 0x64, 0x46, 0x5b, 0x92, 0x98, 0xa8, 0x45, 0xa7, 0xfe, 0xc1, 0xbb]; + let frame = encode(&CborValue::Map(vec![(text("job_id"), CborValue::Tag(37, Box::new(CborValue::Bytes(uuid_bytes)))), (text("current_status"), text("completed")), (text("percentage_complete"), CborValue::Integer(100.into()))])); + + let event = MidjourneyWsEvent::from_cbor_frame(&frame).unwrap(); + let MidjourneyWsEvent::Completed { job_id, percentage_complete } = event else { + panic!("expected Completed, got {event:?}"); + }; + assert_eq!(job_id, "7da8da16-6964-465b-9298-a845a7fec1bb"); + assert_eq!(percentage_complete, 100.0); + } + + #[test] + fn decodes_progress_with_images() { + let jpeg = vec![0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10]; + let frame = encode(&CborValue::Map(vec![(text("current_status"), text("running")), (text("percentage_complete"), CborValue::Float(5.0)), (text("job_id"), text("job-1")), (text("user_id"), text("user-1")), (text("files"), CborValue::Map(vec![(text("0_step_1.jpeg"), CborValue::Bytes(jpeg.clone())), (text("1_step_1.jpeg"), CborValue::Bytes(jpeg.clone()))]))])); + + let event = MidjourneyWsEvent::from_cbor_frame(&frame).unwrap(); + let MidjourneyWsEvent::Progress(progress) = event else { + panic!("expected Progress, got {event:?}"); + }; + assert_eq!(progress.current_status, "running"); + assert_eq!(progress.percentage_complete, 5.0); + assert_eq!(progress.images.len(), 2); + assert_eq!(progress.images[0].image_index, 0); + assert_eq!(progress.images[1].image_index, 1); + } + + #[test] + fn non_map_frame_errors() { + let frame = encode(&text("nope")); + assert!(MidjourneyWsEvent::from_cbor_frame(&frame).is_err()); + } +} diff --git a/crates/api_clients/midjourney_client/src/client/websocket/mod.rs b/crates/api_clients/midjourney_client/src/client/websocket/mod.rs new file mode 100644 index 0000000000..6383f36be6 --- /dev/null +++ b/crates/api_clients/midjourney_client/src/client/websocket/mod.rs @@ -0,0 +1,15 @@ +//! Midjourney websocket client. +//! +//! Midjourney streams job progress (and in-progress preview images) over a +//! websocket at `wss://ws.midjourney.com/ws`. Client-to-server frames are JSON +//! text; server-to-client frames are CBOR binary. +//! +//! Usage: open the socket with [`open_midjourney_websocket`], submit a job over +//! HTTP (see the `submit_job` endpoint), then read previews with +//! [`MidjourneyWebSocket::stream_job`] or the raw [`MidjourneyWebSocket::events`] +//! receiver. + +pub mod job_progress; +pub mod midjourney_websocket; +pub mod midjourney_ws_event; +pub mod open_midjourney_websocket; diff --git a/crates/api_clients/midjourney_client/src/client/websocket/open_midjourney_websocket.rs b/crates/api_clients/midjourney_client/src/client/websocket/open_midjourney_websocket.rs new file mode 100644 index 0000000000..e373a8a433 --- /dev/null +++ b/crates/api_clients/midjourney_client/src/client/websocket/open_midjourney_websocket.rs @@ -0,0 +1,132 @@ +use crate::client::midjourney_hostname::MidjourneyHostname; +use crate::client::websocket::midjourney_websocket::MidjourneyWebSocket; +use crate::client::websocket::midjourney_ws_event::MidjourneyWsEvent; +use crate::credentials::midjourney_user_id::MidjourneyUserId; +use crate::error::midjourney_client_error::MidjourneyClientError; +use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; +use std::time::Duration; +use tokio::sync::broadcast; +use tokio::time::timeout; + +/// Midjourney's websocket gateway. The host is fixed regardless of which +/// `www`/custom hostname the HTTP API is reached through. +const WEBSOCKET_URL_BASE: &str = "wss://ws.midjourney.com/ws"; + +/// Protocol version pinned by the web client's `?v=` query parameter. +const WEBSOCKET_PROTOCOL_VERSION: &str = "5"; + +/// How long to wait for the `user_success` reply before giving up. +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(15); + +pub struct OpenMidjourneyWebSocketRequest<'a> { + /// The short-lived websocket token from the index page's `initialProps` + /// (see `get_user_info`, which reads `websocketToken`). + pub websocket_token: &'a str, + + /// The Midjourney user id this connection is for. + pub user_id: MidjourneyUserId, + + /// Used only for the `Origin` header; the socket host is always + /// `ws.midjourney.com`. Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + + /// Defaults to [`BrowserProfile::default`] if absent. Should match the + /// browser whose cookies obtained the `websocket_token`. + pub browser: Option, +} + +/// Open a Midjourney websocket, complete the `subscribe_to_user` handshake, +/// and return a live [`MidjourneyWebSocket`] handle. +/// +/// The returned handle is already bound to the user and ready for +/// [`MidjourneyWebSocket::subscribe_to_job`] / [`MidjourneyWebSocket::stream_job`]. +/// No cookies are needed — the websocket authenticates via `websocket_token`. +pub async fn open_midjourney_websocket(req: OpenMidjourneyWebSocketRequest<'_>) -> Result { + let client = req.browser.clone().unwrap_or_default().build_client().map_err(MidjourneyClientError::WreqError)?; + + let default_hostname = MidjourneyHostname::Standard; + let hostname = req.hostname.unwrap_or(&default_hostname); + let origin = format!("https://{}", hostname.as_str()); + let url = format!("{}?token={}&v={}", WEBSOCKET_URL_BASE, req.websocket_token, WEBSOCKET_PROTOCOL_VERSION,); + + let response = client.websocket(url).header("Origin", origin).send().await.map_err(|_| MidjourneyClientError::WebSocketUpgrade("Websocket upgrade failed".to_string()))?; + + let websocket = response.into_websocket().await.map_err(|_| MidjourneyClientError::WebSocketUpgrade("Websocket upgrade failed".to_string()))?; + + let handle = MidjourneyWebSocket::spawn(req.user_id, websocket); + + // Subscribe to the event stream BEFORE sending, so we cannot miss the reply. + let mut events = handle.events(); + handle.subscribe_to_user()?; + + match timeout(HANDSHAKE_TIMEOUT, await_user_success(&mut events)).await { + Ok(Ok(())) => Ok(handle), + Ok(Err(err)) => Err(err.into()), + Err(_) => Err(MidjourneyClientError::WebSocketHandshakeTimeout.into()), + } +} + +async fn await_user_success(events: &mut broadcast::Receiver>) -> Result<(), MidjourneyClientError> { + loop { + match events.recv().await { + Ok(event) => { + if matches!(event.as_ref(), MidjourneyWsEvent::UserSuccess { .. }) { + return Ok(()); + } + }, + Err(broadcast::error::RecvError::Lagged(_)) => continue, + Err(broadcast::error::RecvError::Closed) => { + return Err(MidjourneyClientError::WebSocketClosed); + }, + } + } +} + +#[cfg(test)] +mod tests { + use crate::client::websocket::midjourney_ws_event::MidjourneyWsEvent; + use crate::client::websocket::open_midjourney_websocket::{open_midjourney_websocket, OpenMidjourneyWebSocketRequest}; + use crate::endpoints::submit_job::{submit_job, SubmitJobArgs, SubmitJobRequest}; + use crate::recipes::channel_id::ChannelId; + use crate::recipes::get_user_info::{get_user_info, GetUserInfoArgs}; + use errors::AnyhowResult; + use filesys::read_to_trimmed_string::read_to_trimmed_string; + + // Live end-to-end test. Requires a valid cookie header on disk. Run with: + // cargo test -p midjourney_client -- --ignored --exact \ + // client::websocket::open_midjourney_websocket::tests::live_stream_job + #[ignore] + #[tokio::test] + async fn live_stream_job() -> AnyhowResult<()> { + let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?; + + // 1. Read user id + websocket token from the index page. + let user_info = get_user_info(GetUserInfoArgs { cookie_header: &cookie_header, hostname: None, browser: None }).await?; + + let user_id = user_info.user_id.expect("user id"); + let websocket_token = user_info.websocket_token.expect("websocket token"); + + // 2. Open the websocket (does the subscribe_to_user handshake). + let ws = open_midjourney_websocket(OpenMidjourneyWebSocketRequest { websocket_token: &websocket_token, user_id: user_id.clone(), hostname: None, browser: None }).await?; + + // 3. Submit a job over HTTP. + let channel_id = ChannelId::UserId(user_id.clone()).to_string(); + let submit = submit_job(SubmitJobArgs { request: SubmitJobRequest { prompt: "pirate ship in the city --v 8.2", channel_id: &channel_id }, cookie_header: &cookie_header, hostname: None, browser: None }).await?; + let job_id = submit.maybe_job_id.expect("job id"); + + // 4. Stream previews until completion. + let mut preview_frames = 0; + ws.stream_job(&job_id, |event| { + if let MidjourneyWsEvent::Progress(progress) = event { + if !progress.images.is_empty() { + preview_frames += 1; + } + } + }) + .await?; + + assert!(preview_frames > 0, "expected at least one preview frame"); + Ok(()) + } +} diff --git a/crates/api_clients/midjourney_client/src/credentials/cookie_store_has_auth_cookies.rs b/crates/api_clients/midjourney_client/src/credentials/cookie_store_has_auth_cookies.rs index 89827f00f9..007cd94a19 100644 --- a/crates/api_clients/midjourney_client/src/credentials/cookie_store_has_auth_cookies.rs +++ b/crates/api_clients/midjourney_client/src/credentials/cookie_store_has_auth_cookies.rs @@ -1,5 +1,5 @@ use crate::credentials::cookie_names::{AUTH_COOKIE_NAME_I, AUTH_COOKIE_NAME_R}; -use cookie_store::cookie_store::CookieStore; +use cookie_store_wrapper::cookie_store::CookieStore; /// Returns true if the given cookie store has the necessary authentication cookies. pub fn cookie_store_has_auth_cookies(cookie_store: &CookieStore) -> bool { diff --git a/crates/api_clients/midjourney_client/src/endpoints/get_index_page_html.rs b/crates/api_clients/midjourney_client/src/endpoints/get_index_page_html.rs index 90949efa9d..d5f1f9977b 100644 --- a/crates/api_clients/midjourney_client/src/endpoints/get_index_page_html.rs +++ b/crates/api_clients/midjourney_client/src/endpoints/get_index_page_html.rs @@ -2,54 +2,54 @@ use crate::client::midjourney_hostname::MidjourneyHostname; use crate::error::midjourney_api_error::MidjourneyApiError; use crate::error::midjourney_client_error::MidjourneyClientError; use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors; use log::error; -use serde::Deserialize; -use wreq::Client; -use wreq_util::Emulation; -/// This grabs the pre-rendered HTML page, which contains the websocket token and other details. -pub struct GetIndexPageRequest { - pub hostname: MidjourneyHostname, - pub cookie_header: String, +/// Fetches the pre-rendered index HTML (which contains the websocket token and +/// other details). Has no semantic parameters — only transport concerns. +pub struct GetIndexPageArgs<'a> { + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, } -pub async fn get_index_page_html(req: GetIndexPageRequest) -> Result { - let url = format!("https://{}/", req.hostname.as_str()); +pub async fn get_index_page_html(args: GetIndexPageArgs<'_>) -> Result { + let default_hostname = MidjourneyHostname::Standard; + let hostname = args.hostname.unwrap_or(&default_hostname); - let client = Client::builder() - .emulation(Emulation::Firefox139) - .build() - .map_err(|err| MidjourneyClientError::WreqError(err))?; + let url = format!("https://{}/", hostname.as_str()); - let cookie_header = req.cookie_header.trim(); + let client = args.browser.clone().unwrap_or_default() + .build_client() + .map_err(MidjourneyClientError::WreqError)?; + + let cookie_header = args.cookie_header.trim(); if cookie_header.len() < 20 { - error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header); + error!("Cookie header is too short (len: {})", cookie_header.len()); return Err(MidjourneyClientError::CookieTooShort.into()); } - // NB: missing headers that were in the browser request: - // -H 'sec-ch-ua-platform: "macOS"' \ - // -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \ - - let mut http_request = client.get(url) + // NB: Browser-identity headers (user-agent, sec-ch-ua*, accept-encoding) come + // from the emulation on the client; only request-context headers are set here. + let http_request = client.get(url) .header("cookie", cookie_header) .header("Referrer-Policy", "origin-when-cross-origin") .header("accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7") .header("accept-language", "en-US,en;q=0.8") - .header("content-type", "application/json") .header("priority", "u=0, i") - .header("sec-ch-ua-mobile", "?0") .header("sec-fetch-dest", "document") .header("sec-fetch-mode", "navigate") .header("sec-fetch-site", "none") .header("sec-fetch-user", "?1") .header("upgrade-insecure-requests", "1"); - let http_request = http_request + let http_request = http_request .build() - .map_err(|err| MidjourneyClientError::WreqError(err))?; + .map_err(MidjourneyClientError::WreqError)?; let response = client.execute(http_request) .await @@ -65,25 +65,25 @@ pub async fn get_index_page_html(req: GetIndexPageRequest) -> Result AnyhowResult<()> { let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?; - let result = get_index_page_html(GetIndexPageRequest { - cookie_header, - hostname: MidjourneyHostname::Standard, + let result = get_index_page_html(GetIndexPageArgs { + cookie_header: &cookie_header, + hostname: None, + browser: None, }).await?; println!("Response: {:?}\n\n", result); diff --git a/crates/api_clients/midjourney_client/src/endpoints/imagine.rs b/crates/api_clients/midjourney_client/src/endpoints/imagine.rs index 35ea862e20..0e9d739925 100644 --- a/crates/api_clients/midjourney_client/src/endpoints/imagine.rs +++ b/crates/api_clients/midjourney_client/src/endpoints/imagine.rs @@ -3,22 +3,29 @@ use crate::credentials::midjourney_user_id::MidjourneyUserId; use crate::error::midjourney_api_error::MidjourneyApiError; use crate::error::midjourney_client_error::MidjourneyClientError; use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors; use log::error; use serde::Deserialize; -use wreq::Client; -use wreq_util::Emulation; const DEFAULT_PAGE_SIZE : u64 = 1000; -/// This returns Midjourney's job/media list +/// The semantic parameters for listing Midjourney's job/media feed. pub struct ImagineRequest<'a> { - pub hostname: MidjourneyHostname, - pub cookie_header: String, pub user_id: &'a MidjourneyUserId, pub page_size: Option, } +/// An imagine (job/media list) request plus its transport concerns. +pub struct ImagineArgs<'a> { + pub request: ImagineRequest<'a>, + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, +} + #[derive(Debug, Clone)] pub struct ImagineResponse { pub items: Vec, @@ -71,43 +78,41 @@ impl MidjourneyJobType { } } -pub async fn imagine(req: ImagineRequest<'_>) -> Result { - let cookie_header = req.cookie_header.trim(); +pub async fn imagine(args: ImagineArgs<'_>) -> Result { + let cookie_header = args.cookie_header.trim(); if cookie_header.len() < 20 { - error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header); + error!("Cookie header is too short (len: {})", cookie_header.len()); return Err(MidjourneyClientError::CookieTooShort.into()); } - let page_size = req.page_size.unwrap_or(DEFAULT_PAGE_SIZE); + let default_hostname = MidjourneyHostname::Standard; + let hostname = args.hostname.unwrap_or(&default_hostname); - let referer = format!("https://{}/imagine", req.hostname.as_str()); + let page_size = args.request.page_size.unwrap_or(DEFAULT_PAGE_SIZE); + + let referer = format!("https://{}/imagine", hostname.as_str()); let url = format!("https://{}/api/imagine?user_id={}&page_size={}", - req.hostname.as_str(), - req.user_id.as_str(), + hostname.as_str(), + args.request.user_id.as_str(), page_size, ); - let client = Client::builder() - .emulation(Emulation::Firefox139) - .build() - .map_err(|err| MidjourneyClientError::WreqError(err))?; + let client = args.browser.clone().unwrap_or_default() + .build_client() + .map_err(MidjourneyClientError::WreqError)?; - // NB: missing headers that were in the browser request: - // -H 'sec-ch-ua-platform: "macOS"' \ - // -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \ - - let mut http_request = client.get(url) + // NB: Browser-identity headers (user-agent, sec-ch-ua*, accept-encoding) come + // from the emulation on the client; only request-context headers are set here. + let http_request = client.get(url) .header("cookie", cookie_header) .header("Referer", &referer) .header("Referrer-Policy", "origin-when-cross-origin") .header("accept", "*/*") - .header("accept-encoding", "gzip, deflate, br, zstd") .header("accept-language", "en-US,en;q=0.8") .header("content-type", "application/json") .header("priority", "u=1, i") - .header("sec-ch-ua-mobile", "?0") .header("sec-fetch-dest", "empty") .header("sec-fetch-mode", "cors") .header("sec-fetch-site", "same-origin") @@ -201,7 +206,7 @@ pub async fn imagine(req: ImagineRequest<'_>) -> Result) -> Result { + /// A client-generated session id (UUID) grouping prompts in one sitting. + pub session_id: &'a str, + + /// Monotonic index of this prompt within the session, starting at 0. + pub sequence_index: u32, + + /// Epoch milliseconds when the prompt was captured. + pub captured_at_ms: u64, + + /// The prompt text WITHOUT the trailing `--` parameters. + pub raw_text: &'a str, + + /// The model version string, e.g. `"8.2"`. + pub version: &'a str, +} + +/// A prompt-session-log request plus its transport concerns. +pub struct PromptSessionLogArgs<'a> { + pub request: PromptSessionLogRequest<'a>, + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, +} + +#[derive(Debug, Clone)] +pub struct PromptSessionLogResponse { + /// The `status` field from the response body (e.g. `"success"`). + pub status: Option, +} + +pub async fn prompt_session_log(args: PromptSessionLogArgs<'_>) -> Result { + let cookie_header = args.cookie_header.trim(); + if cookie_header.len() < 20 { + error!("Cookie header is too short (len: {})", cookie_header.len()); + return Err(MidjourneyClientError::CookieTooShort.into()); + } + + let default_hostname = MidjourneyHostname::Standard; + let hostname = args.hostname.unwrap_or(&default_hostname); + + let client = args.browser.clone().unwrap_or_default().build_client().map_err(MidjourneyClientError::WreqError)?; + + let referer = format!("https://{}/imagine", hostname.as_str()); + let url = format!("https://{}/api/prompt-session-log", hostname.as_str()); + + #[derive(Serialize)] + struct RawRequest<'a> { + session_id: &'a str, + sequence_index: u32, + captured_at_ms: u64, + source: &'a str, + raw_text: &'a str, + route: &'a str, + in_conversation_mode: bool, + version: &'a str, + video_mode: bool, + image_prompt_count: u32, + style_ref_count: u32, + character_ref_count: u32, + depth_ref_count: u32, + starting_frame_count: u32, + ending_frame_count: u32, + openai_item_id: Option, + } + + let body = RawRequest { session_id: args.request.session_id, sequence_index: args.request.sequence_index, captured_at_ms: args.request.captured_at_ms, source: "imagine_text", raw_text: args.request.raw_text, route: "/imagine", in_conversation_mode: false, version: args.request.version, video_mode: false, image_prompt_count: 0, style_ref_count: 0, character_ref_count: 0, depth_ref_count: 0, starting_frame_count: 0, ending_frame_count: 0, openai_item_id: None }; + + let http_request = client.post(url).header("cookie", cookie_header).header("Referer", &referer).header("Referrer-Policy", "origin-when-cross-origin").header("accept", "*/*").header("accept-language", "en-US,en;q=0.8").header("content-type", "application/json").header("priority", "u=1, i").header("sec-fetch-dest", "empty").header("sec-fetch-mode", "cors").header("sec-fetch-site", "same-origin").header("x-csrf-protection", "1").json(&body).build().map_err(MidjourneyClientError::WreqError)?; + + let response = client.execute(http_request).await.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?; + + let status = response.status(); + let response_body = response.text().await.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?; + + if !status.is_success() { + if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) { + return Err(MidjourneyApiError::CloudflareError(err).into()); + } + return Err(MidjourneyApiError::UnknownHttpFailure { status_code: status.as_u16(), body: response_body }.into()); + } + + // {"status":"success"} + #[derive(Deserialize)] + struct RawResponse { + status: Option, + } + + let response = serde_json::from_str::(&response_body).map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?; + + Ok(PromptSessionLogResponse { status: response.status }) +} diff --git a/crates/api_clients/midjourney_client/src/endpoints/refresh_token.rs b/crates/api_clients/midjourney_client/src/endpoints/refresh_token.rs new file mode 100644 index 0000000000..04ac618bc1 --- /dev/null +++ b/crates/api_clients/midjourney_client/src/endpoints/refresh_token.rs @@ -0,0 +1,96 @@ +use crate::error::midjourney_api_error::MidjourneyApiError; +use crate::error::midjourney_client_error::MidjourneyClientError; +use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; +use serde::Deserialize; + +/// Google Identity Toolkit endpoint that Midjourney's Firebase auth uses to +/// exchange a refresh token for a fresh ID token. +const SECURE_TOKEN_URL: &str = "https://securetoken.googleapis.com/v1/token"; + +/// The public Firebase Web API key embedded in the Midjourney web app. This is +/// not a secret (it ships in the client bundle); it identifies the Firebase +/// project, and access is still gated by the refresh token. +pub const DEFAULT_FIREBASE_API_KEY: &str = "AIzaSyAjizp68NsH3JGUS0EyLXsChW4fN0A92tM"; + +/// Exchanges a refresh token for a new ID/access token. +/// +/// The refresh token is the value of the `__Host-Midjourney.AuthUserTokenV3_r` +/// cookie. The returned `id_token` is the new value for the +/// `__Host-Midjourney.AuthUserTokenV3_i` cookie. +pub struct RefreshTokenRequest<'a> { + pub refresh_token: &'a str, + + /// Firebase Web API key. Pass `DEFAULT_FIREBASE_API_KEY` unless overriding. + pub api_key: &'a str, +} + +/// A refresh-token request plus its transport concerns. This endpoint hits +/// Google's securetoken service (not Midjourney), so it has no cookie or +/// hostname — only an optional browser profile. +pub struct RefreshTokenArgs<'a> { + pub request: RefreshTokenRequest<'a>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, +} + +#[derive(Debug, Clone)] +pub struct RefreshTokenResponse { + /// The new access token (same value as `id_token`). + pub access_token: String, + + /// The new Firebase ID token — the fresh `AuthUserTokenV3_i` cookie value. + pub id_token: String, + + /// The refresh token to use next time (may be rotated). + pub refresh_token: String, + + /// Seconds until `access_token` expires (as returned, a stringified number). + pub expires_in: Option, + + pub token_type: Option, + pub user_id: Option, + pub project_id: Option, +} + +pub async fn refresh_token(args: RefreshTokenArgs<'_>) -> Result { + let client = args.browser.clone().unwrap_or_default().build_client().map_err(MidjourneyClientError::WreqError)?; + + let url = format!("{}?key={}", SECURE_TOKEN_URL, args.request.api_key); + + let form = [("grant_type", "refresh_token"), ("refresh_token", args.request.refresh_token)]; + + let http_request = client.post(url).header("Origin", "https://www.midjourney.com").header("Referer", "https://www.midjourney.com/").header("accept", "*/*").header("accept-language", "en-US,en;q=0.8").header("content-type", "application/x-www-form-urlencoded").form(&form).build().map_err(MidjourneyClientError::WreqError)?; + + let response = client.execute(http_request).await.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?; + + let status = response.status(); + let response_body = response.text().await.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?; + + if !status.is_success() { + return Err( + match status.as_u16() { + 400 => MidjourneyApiError::InvalidRequest(response_body), + 401 => MidjourneyApiError::Unauthorized(response_body), + 403 => MidjourneyApiError::Forbidden(response_body), + code => MidjourneyApiError::UnknownHttpFailure { status_code: code, body: response_body }, + } + .into(), + ); + } + + #[derive(Deserialize)] + struct RawResponse { + access_token: String, + id_token: String, + refresh_token: String, + expires_in: Option, + token_type: Option, + user_id: Option, + project_id: Option, + } + + let response = serde_json::from_str::(&response_body).map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?; + + Ok(RefreshTokenResponse { access_token: response.access_token, id_token: response.id_token, refresh_token: response.refresh_token, expires_in: response.expires_in, token_type: response.token_type, user_id: response.user_id, project_id: response.project_id }) +} diff --git a/crates/api_clients/midjourney_client/src/endpoints/storage_list.rs b/crates/api_clients/midjourney_client/src/endpoints/storage_list.rs index 9fb0ad9f43..7a5903357a 100644 --- a/crates/api_clients/midjourney_client/src/endpoints/storage_list.rs +++ b/crates/api_clients/midjourney_client/src/endpoints/storage_list.rs @@ -2,16 +2,19 @@ use crate::client::midjourney_hostname::MidjourneyHostname; use crate::error::midjourney_api_error::MidjourneyApiError; use crate::error::midjourney_client_error::MidjourneyClientError; use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors; use log::error; use serde::Deserialize; -use wreq::Client; -use wreq_util::Emulation; -/// This endpoint returns the user id -pub struct GetStorageListRequest { - pub hostname: MidjourneyHostname, - pub cookie_header: String, +/// Lists the user's storage (the paths contain the user id). Has no semantic +/// parameters — only transport concerns. +pub struct StorageListArgs<'a> { + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, } #[derive(Debug)] @@ -22,29 +25,28 @@ pub struct StorageItem { // NB: Other fields aren't necessary for us right now. } -pub async fn storage_list(req: GetStorageListRequest) -> Result, MidjourneyError> { +pub async fn storage_list(args: StorageListArgs<'_>) -> Result, MidjourneyError> { + let default_hostname = MidjourneyHostname::Standard; + let hostname = args.hostname.unwrap_or(&default_hostname); - let referer = format!("https://{}", req.hostname.as_str()); + let referer = format!("https://{}", hostname.as_str()); - let url = format!("https://{}/api/storage-list", req.hostname.as_str()); + let url = format!("https://{}/api/storage-list", hostname.as_str()); - let client = Client::builder() - .emulation(Emulation::Firefox139) - .build() - .map_err(|err| MidjourneyClientError::WreqError(err))?; + let client = args.browser.clone().unwrap_or_default() + .build_client() + .map_err(MidjourneyClientError::WreqError)?; - let cookie_header = req.cookie_header.trim(); + let cookie_header = args.cookie_header.trim(); if cookie_header.len() < 20 { - error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header); + error!("Cookie header is too short (len: {})", cookie_header.len()); return Err(MidjourneyClientError::CookieTooShort.into()); } - // NB: missing headers that were in the browser request: - // -H 'sec-ch-ua-platform: "macOS"' \ - // -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \ - - let mut http_request = client.get(url) + // NB: Browser-identity headers come from the emulation on the client; only + // request-context headers are set here. + let http_request = client.get(url) .header("cookie", cookie_header) .header("Referer", &referer) .header("Referrer-Policy", "origin-when-cross-origin") @@ -52,7 +54,6 @@ pub async fn storage_list(req: GetStorageListRequest) -> Result .header("accept-language", "en-US,en;q=0.8") .header("content-type", "application/json") .header("priority", "u=1, i") - .header("sec-ch-ua-mobile", "?0") .header("sec-fetch-dest", "empty") .header("sec-fetch-mode", "cors") .header("sec-fetch-site", "same-origin") @@ -105,7 +106,7 @@ pub async fn storage_list(req: GetStorageListRequest) -> Result } let response : Vec = serde_json::from_str(response_body) - .map_err(|err| MidjourneyApiError::DeserializationError(err))?; + .map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?; let items = response .into_iter() @@ -119,19 +120,19 @@ pub async fn storage_list(req: GetStorageListRequest) -> Result #[cfg(test)] mod tests { - use crate::client::midjourney_hostname::MidjourneyHostname; use errors::AnyhowResult; use filesys::read_to_trimmed_string::read_to_trimmed_string; - use crate::endpoints::storage_list::{storage_list, GetStorageListRequest}; + use crate::endpoints::storage_list::{storage_list, StorageListArgs}; #[ignore] #[tokio::test] async fn test() -> AnyhowResult<()> { let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?; - let result = storage_list(GetStorageListRequest { - cookie_header, - hostname: MidjourneyHostname::Standard, + let result = storage_list(StorageListArgs { + cookie_header: &cookie_header, + hostname: None, + browser: None, }).await?; println!("Response: {:?}\n\n", result); diff --git a/crates/api_clients/midjourney_client/src/endpoints/submit_job.rs b/crates/api_clients/midjourney_client/src/endpoints/submit_job.rs index c1240ee20d..59e5bdef68 100644 --- a/crates/api_clients/midjourney_client/src/endpoints/submit_job.rs +++ b/crates/api_clients/midjourney_client/src/endpoints/submit_job.rs @@ -5,51 +5,101 @@ use log::{error, warn}; use serde::{Deserialize, Serialize}; use crate::client::midjourney_hostname::MidjourneyHostname; +use browser_emulation::browser_profile::BrowserProfile; use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors; -use wreq::Client; -use wreq_util::Emulation; +/// The semantic parameters of a submit-jobs request. pub struct SubmitJobRequest<'a> { pub prompt: &'a str, pub channel_id: &'a str, - pub hostname: MidjourneyHostname, - pub cookie_header: String, +} + +/// A submit-jobs request plus its transport concerns. +pub struct SubmitJobArgs<'a> { + pub request: SubmitJobRequest<'a>, + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, } #[derive(Debug, Clone)] pub struct SubmitJobResponse { /// On success, the job ID is returned. pub maybe_job_id: Option, - + /// On error, we have a list of error messages. pub maybe_errors: Option>, } +impl SubmitJobResponse { + /// The classified errors returned by the submit, if any. + pub fn errors(&self) -> &[SubmitJobError] { + self.maybe_errors.as_deref().unwrap_or(&[]) + } + + /// Whether the submit was rejected because the account has no active paid + /// Midjourney subscription (free trials disabled). First-class so callers + /// can branch on it (e.g. prompt the user to subscribe). + pub fn is_subscription_required(&self) -> bool { + self.errors() + .iter() + .any(|error| error.error_type == MidjourneySubmitErrorType::SubscriptionRequired) + } +} + #[derive(Debug, Clone)] pub struct SubmitJobError { - pub error_type: Option, + pub error_type: MidjourneySubmitErrorType, pub message: Option, } -pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result { - let client = Client::builder() - .emulation(Emulation::Firefox139) - .build() - .map_err(|err| MidjourneyClientError::WreqError(err))?; +/// A classified Midjourney submit-jobs failure `type`. Unknown values are +/// preserved via [`MidjourneySubmitErrorType::Other`] so matching stays +/// exhaustive-safe. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum MidjourneySubmitErrorType { + /// `subscription_required`: the account needs an active paid Midjourney plan. + SubscriptionRequired, - let referer = format!("https://{}", req.hostname.as_str()); + /// Any other failure type Midjourney returned (carries the raw string). + Other(String), +} + +impl MidjourneySubmitErrorType { + pub fn from_api_str(value: &str) -> Self { + match value { + "subscription_required" => Self::SubscriptionRequired, + other => Self::Other(other.to_string()), + } + } +} + +pub async fn submit_job(args: SubmitJobArgs<'_>) -> Result { + let default_hostname = MidjourneyHostname::Standard; + let hostname = args.hostname.unwrap_or(&default_hostname); + + let client = args.browser.clone().unwrap_or_default() + .build_client() + .map_err(MidjourneyClientError::WreqError)?; + + let referer = format!("https://{}", hostname.as_str()); // NB: Other recent clients use /api/app/submit-jobs, but this appears invalid. - let url = format!("https://{}/api/submit-jobs", req.hostname.as_str()); + let url = format!("https://{}/api/submit-jobs", hostname.as_str()); - let cookie_header = req.cookie_header.trim(); + let cookie_header = args.cookie_header.trim(); if cookie_header.len() < 20 { - error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header); + error!("Cookie header is too short (len: {})", cookie_header.len()); return Err(MidjourneyClientError::CookieTooShort.into()); } - let mut http_request = client.post(url) + // NB: Browser-identity headers (user-agent, sec-ch-ua*, accept-encoding) are + // set coherently by the emulation on the client; here we set only the + // request-context headers. + let http_request = client.post(url) .header("cookie", cookie_header) .header("Referer", &referer) .header("Referrer-Policy", "origin-when-cross-origin") @@ -57,7 +107,6 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result) -> Result, imagePrompts: u8, imageReferences: u8, characterReferences: u8, depthReferences: u8, - lightboxOpen: String, + lightboxOpen: Option, } #[derive(Serialize)] @@ -94,16 +145,17 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result) -> Result) -> Result(response_body) - .map_err(|err| MidjourneyApiError::DeserializationError(err))?; + .map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?; let maybe_job_id = response.success .get(0) @@ -185,9 +244,9 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result) -> Result { + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, } #[derive(Debug)] @@ -30,29 +32,28 @@ pub struct Provider { pub is_linked: bool, } -pub async fn user_connections(req: GetUserConnectionsRequest) -> Result, MidjourneyError> { +pub async fn user_connections(args: UserConnectionsArgs<'_>) -> Result, MidjourneyError> { + let default_hostname = MidjourneyHostname::Standard; + let hostname = args.hostname.unwrap_or(&default_hostname); - let referer = format!("https://{}", req.hostname.as_str()); + let referer = format!("https://{}", hostname.as_str()); - let url = format!("https://{}/api/user-connections", req.hostname.as_str()); + let url = format!("https://{}/api/user-connections", hostname.as_str()); - let client = Client::builder() - .emulation(Emulation::Firefox139) - .build() - .map_err(|err| MidjourneyClientError::WreqError(err))?; + let client = args.browser.clone().unwrap_or_default() + .build_client() + .map_err(MidjourneyClientError::WreqError)?; - let cookie_header = req.cookie_header.trim(); + let cookie_header = args.cookie_header.trim(); if cookie_header.len() < 20 { - error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header); + error!("Cookie header is too short (len: {})", cookie_header.len()); return Err(MidjourneyClientError::CookieTooShort.into()); } - // NB: missing headers that were in the browser request: - // -H 'sec-ch-ua-platform: "macOS"' \ - // -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \ - - let mut http_request = client.get(url) + // NB: Browser-identity headers come from the emulation on the client; only + // request-context headers are set here. + let http_request = client.get(url) .header("cookie", cookie_header) .header("Referer", &referer) .header("Referrer-Policy", "origin-when-cross-origin") @@ -60,7 +61,6 @@ pub async fn user_connections(req: GetUserConnectionsRequest) -> Result Result = serde_json::from_str(response_body) - .map_err(|err| MidjourneyApiError::DeserializationError(err))?; + .map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?; let providers = response .into_iter() @@ -149,8 +149,7 @@ pub async fn user_connections(req: GetUserConnectionsRequest) -> Result AnyhowResult<()> { let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?; - let result = user_connections(GetUserConnectionsRequest { - cookie_header, - hostname: MidjourneyHostname::Standard, + let result = user_connections(UserConnectionsArgs { + cookie_header: &cookie_header, + hostname: None, + browser: None, }).await?; println!("Response: {:?}\n\n", result); diff --git a/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs b/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs index 59081f7d7e..13b3d2e6d2 100644 --- a/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs +++ b/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs @@ -47,8 +47,11 @@ pub enum MidjourneyApiError { /// Cloudflare errors. CloudflareError(CloudflareError), - /// A deserialization error with the response. - DeserializationError(serde_json::Error), + /// A response could not be decoded. The payload may contain credentials. + DeserializationError { + source: serde_json::Error, + body: String, + }, /// The request timed out. Timeout(String), @@ -84,7 +87,13 @@ impl Display for MidjourneyApiError { write!(f, "Unknown HTTP failure; status code: {}; body: {}", status_code, body), // Deserialization errors // Server response handling errors - Self::DeserializationError(error) => write!(f, "Deserialization error: {}", error), + Self::DeserializationError { source, body } => write!( + f, + "Deserialization error: {} | raw response body ({} bytes): {}", + source, + body.len(), + body, + ), // Network errors Self::Timeout(msg) => write!(f, "Timeout: {}", msg), Self::NetworkError(msg) => write!(f, "Network error: {}", msg), @@ -98,9 +107,9 @@ impl Display for MidjourneyApiError { } } -impl From for MidjourneyApiError { - fn from(error: serde_json::Error) -> Self { - Self::DeserializationError(error) +impl MidjourneyApiError { + pub fn deserialization(source: serde_json::Error, _body: &str) -> Self { + Self::DeserializationError { source, body: "".to_string() } } } 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 index de74e8e1cd..2682926773 100644 --- a/crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs +++ b/crates/api_clients/midjourney_client/src/error/midjourney_client_error.rs @@ -27,6 +27,16 @@ pub enum MidjourneyClientError { /// A serialization error with the request. SerializationError(serde_json::Error), + + /// The websocket upgrade handshake failed. + WebSocketUpgrade(String), + + /// Tried to send on a websocket whose background task has stopped. + WebSocketClosed, + + /// The websocket opened but the `subscribe_to_user` handshake did not + /// complete within the allotted time. + WebSocketHandshakeTimeout, } impl Error for MidjourneyClientError {} @@ -42,6 +52,9 @@ impl Display for MidjourneyClientError { //Self::ReqwestError(err) => write!(f, "Reqwest client error: {}", err), Self::WreqError(err) => write!(f, "Wreq client error: {}", err), 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::WebSocketHandshakeTimeout => write!(f, "Timed out waiting for the websocket user handshake."), } } } diff --git a/crates/api_clients/midjourney_client/src/recipes/get_midjourney_user_id_from_storage.rs b/crates/api_clients/midjourney_client/src/recipes/get_midjourney_user_id_from_storage.rs index 1c0a112452..25182173f9 100644 --- a/crates/api_clients/midjourney_client/src/recipes/get_midjourney_user_id_from_storage.rs +++ b/crates/api_clients/midjourney_client/src/recipes/get_midjourney_user_id_from_storage.rs @@ -1,18 +1,27 @@ use crate::client::midjourney_hostname::MidjourneyHostname; use crate::credentials::midjourney_user_id::MidjourneyUserId; -use crate::endpoints::storage_list::{storage_list, GetStorageListRequest}; +use crate::endpoints::storage_list::{storage_list, StorageListArgs}; use crate::error::midjourney_api_error::MidjourneyApiError; use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; -pub struct GetMidjourneyUserIdRequest { - pub hostname: MidjourneyHostname, - pub cookie_header: String, +/// Recovers the user id from the storage list. Has no semantic parameters — +/// only transport concerns. +pub struct GetMidjourneyUserIdArgs<'a> { + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, } -pub async fn get_midjourney_user_id_from_storage(req: GetMidjourneyUserIdRequest) -> Result { - let items = storage_list(GetStorageListRequest { - hostname: req.hostname, - cookie_header: req.cookie_header, +pub async fn get_midjourney_user_id_from_storage( + args: GetMidjourneyUserIdArgs<'_>, +) -> Result { + let items = storage_list(StorageListArgs { + cookie_header: args.cookie_header, + hostname: args.hostname, + browser: args.browser, }).await?; let user_id = items @@ -27,8 +36,7 @@ pub async fn get_midjourney_user_id_from_storage(req: GetMidjourneyUserIdRequest #[cfg(test)] mod tests { - use crate::client::midjourney_hostname::MidjourneyHostname; - use crate::recipes::get_midjourney_user_id_from_storage::{get_midjourney_user_id_from_storage, GetMidjourneyUserIdRequest}; + use crate::recipes::get_midjourney_user_id_from_storage::{get_midjourney_user_id_from_storage, GetMidjourneyUserIdArgs}; use errors::AnyhowResult; use filesys::read_to_trimmed_string::read_to_trimmed_string; @@ -37,9 +45,10 @@ mod tests { async fn test() -> AnyhowResult<()> { let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?; - let result = get_midjourney_user_id_from_storage(GetMidjourneyUserIdRequest { - cookie_header, - hostname: MidjourneyHostname::Standard, + let result = get_midjourney_user_id_from_storage(GetMidjourneyUserIdArgs { + cookie_header: &cookie_header, + hostname: None, + browser: None, }).await?; println!("Response: {:?}\n\n", result); @@ -49,4 +58,3 @@ mod tests { Ok(()) } } - diff --git a/crates/api_clients/midjourney_client/src/recipes/get_user_info.rs b/crates/api_clients/midjourney_client/src/recipes/get_user_info.rs index f1fee8f084..43c25298ba 100644 --- a/crates/api_clients/midjourney_client/src/recipes/get_user_info.rs +++ b/crates/api_clients/midjourney_client/src/recipes/get_user_info.rs @@ -1,14 +1,20 @@ use crate::client::midjourney_hostname::MidjourneyHostname; use crate::credentials::midjourney_user_id::MidjourneyUserId; -use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageRequest}; +use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageArgs}; use crate::error::midjourney_api_error::MidjourneyApiError; use crate::error::midjourney_error::MidjourneyError; +use browser_emulation::browser_profile::BrowserProfile; use dom_query::Document; use serde::Deserialize; -pub struct GetUserInfoRequest { - pub hostname: MidjourneyHostname, - pub cookie_header: String, +/// Fetches the user's id, email, and websocket token from the index page. Has +/// no semantic parameters — only transport concerns. +pub struct GetUserInfoArgs<'a> { + pub cookie_header: &'a str, + /// Defaults to the standard hostname if absent. + pub hostname: Option<&'a MidjourneyHostname>, + /// Defaults to [`BrowserProfile::default`] if absent. + pub browser: Option, } #[derive(Debug, Clone)] @@ -18,12 +24,25 @@ pub struct GetUserInfoResponse { pub websocket_token: Option, } -pub async fn get_user_info(req: GetUserInfoRequest) -> Result { - let index_html = get_index_page_html(GetIndexPageRequest { - hostname: req.hostname, - cookie_header: req.cookie_header, +pub async fn get_user_info(args: GetUserInfoArgs<'_>) -> Result { + let index_html = get_index_page_html(GetIndexPageArgs { + cookie_header: args.cookie_header, + hostname: args.hostname, + browser: args.browser, }).await?; + // Diagnostics: distinguish a Cloudflare challenge from a genuinely + // logged-out page when the parse below fails. + log::info!( + "Midjourney index page: {} bytes; initialProps={}, initialAuthUser={}, websocketToken={}, turnstile={}, challenge-platform={}", + index_html.len(), + index_html.contains("initialProps"), + index_html.contains("initialAuthUser"), + index_html.contains("websocketToken"), + index_html.contains("turnstile"), + index_html.contains("challenge-platform"), + ); + /*