backport midjourney fixes from ArtCraft-X (#1953)

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",

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" }

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<P: AsRef<Path>>(
   args: UploadImageFromFileArgs<'_, P>,
 ) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
   
+  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<P: AsRef<Path>>(
+  args: UploadImageFromFileArgs<'_, P>,
+  idempotency_token: &str,
+) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
   validate_args(&args)?;
 
   let url = get_route(args.api_host);
@@ -66,7 +75,7 @@ pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
   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 {

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

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,

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;

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<RouterMidjourneyClient>,
   pub(crate) artcraft_client: Option<RouterArtcraftClient>,
   pub(crate) fal_client: Option<RouterFalClient>,
   pub(crate) gmicloud_client: Option<RouterGmiCloudClient>,
@@ -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))

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<RouterMidjourneyClient>,
   artcraft_client: Option<RouterArtcraftClient>,
   fal_client: Option<RouterFalClient>,
   gmicloud_client: Option<RouterGmiCloudClient>,
@@ -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,

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),

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 }
+  }
+}

crates/api_clients/artcraft/artcraft_router/src/errors/artcraft_router_error.rs

@@ -64,6 +64,7 @@ impl From<DownloadError> for ArtcraftRouterError {
 impl From<ProviderError> 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,

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"),

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),

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<ImageGenerationDraftOrRequest, ArtcraftRouterError> {
     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),

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,
+}

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<ImageGenerationCostEstimate, ArtcraftRouterError> {
     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<GenerateImageResponse, ArtcraftRouterError> {
     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()?;

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<ImageGenerationDraftOrRequest, ArtcraftRouterError> {
+  // 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<RouterAspectRatio>, 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<RouterAspectRatio>) -> 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"),
+    }
+  }
+}

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 }
+  }
+}

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;

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<GenerateImageResponse, ArtcraftRouterError> {
+    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::<Vec<String>>().join("; ");
+
+  if detail.is_empty() { "Midjourney rejected the request with no detail".to_string() } else { detail }
+}

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;

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

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;

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<u8>,
+}
+
+/// 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<u8>) -> Option<Self> {
+    // "0_step_5.jpeg" -> index "0", rest "step_5.jpeg"
+    let (index_str, rest) = key.split_once('_')?;
+    let image_index = index_str.parse::<u8>().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::<u32>().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());
+  }
+}

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<WebSocketInner>,
+}
+
+struct WebSocketInner {
+  user_id: MidjourneyUserId,
+  outbound: mpsc::UnboundedSender<Message>,
+  events: broadcast::Sender<Arc<MidjourneyWsEvent>>,
+  connected: Arc<AtomicBool>,
+}
+
+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::<Message>();
+    let (events_tx, _events_rx) = broadcast::channel::<Arc<MidjourneyWsEvent>>(EVENT_CHANNEL_CAPACITY);
+    let connected = Arc::new(AtomicBool::new(true));
+
+    tokio::spawn(run_websocket_loop(websocket, outbound_rx, events_tx.clone(), connected.clone()));
+
+    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<Arc<MidjourneyWsEvent>> {
+    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<F>(&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<Message>, events_tx: broadcast::Sender<Arc<MidjourneyWsEvent>>, connected: Arc<AtomicBool>) {
+  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.");
+}

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<String> },
+
+  /// 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<String> },
+}
+
+/// A single progress frame for a running job.
+#[derive(Clone, Debug)]
+pub struct JobProgress {
+  pub job_id: String,
+  pub maybe_user_id: Option<String>,
+
+  /// 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<JobStepImage>,
+}
+
+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<Self, MidjourneyWsDecodeError> {
+    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<JobStepImage> {
+  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<String> {
+  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<f64> {
+  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<u8> {
+    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());
+  }
+}

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;

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<BrowserProfile>,
+}
+
+/// 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<MidjourneyWebSocket, MidjourneyError> {
+  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<std::sync::Arc<MidjourneyWsEvent>>) -> 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(())
+  }
+}

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 {

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<BrowserProfile>,
 }
 
-pub async fn get_index_page_html(req: GetIndexPageRequest) -> Result<String, MidjourneyError> {
-  let url = format!("https://{}/", req.hostname.as_str());
+pub async fn get_index_page_html(args: GetIndexPageArgs<'_>) -> Result<String, MidjourneyError> {
+  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<String, Mid
       return Err(MidjourneyApiError::CloudflareError(err).into());
     }
   }
-  
+
   Ok(response_body)
 }
 
 #[cfg(test)]
 mod tests {
-  use crate::client::midjourney_hostname::MidjourneyHostname;
+  use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageArgs};
   use errors::AnyhowResult;
   use filesys::read_to_trimmed_string::read_to_trimmed_string;
-  use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageRequest};
 
   #[ignore]
   #[tokio::test]
   async fn test() -> 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);

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<u64>,
 }
 
+/// 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<BrowserProfile>,
+}
+
 #[derive(Debug, Clone)]
 pub struct ImagineResponse {
   pub items: Vec<ImagineItem>,
@@ -71,43 +78,41 @@ impl MidjourneyJobType {
   }
 }
 
-pub async fn imagine(req: ImagineRequest<'_>) -> Result<ImagineResponse, MidjourneyError> {
-  let cookie_header = req.cookie_header.trim();
+pub async fn imagine(args: ImagineArgs<'_>) -> Result<ImagineResponse, MidjourneyError> {
+  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<ImagineResponse, Midjour
   }
 
   let response : RawImagineResponse = serde_json::from_str(&response_body)
-      .map_err(|err| MidjourneyApiError::DeserializationError(err))?;
+      .map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?;
 
   Ok(ImagineResponse {
     cursor: response.cursor.clone(),
@@ -220,9 +225,8 @@ pub async fn imagine(req: ImagineRequest<'_>) -> Result<ImagineResponse, Midjour
 
 #[cfg(test)]
 mod tests {
-  use crate::client::midjourney_hostname::MidjourneyHostname;
   use crate::credentials::midjourney_user_id::MidjourneyUserId;
-  use crate::endpoints::imagine::{imagine, ImagineRequest};
+  use crate::endpoints::imagine::{imagine, ImagineArgs, ImagineRequest};
   use errors::AnyhowResult;
   use filesys::read_to_trimmed_string::read_to_trimmed_string;
 
@@ -233,11 +237,14 @@ mod tests {
     let user_id = read_to_trimmed_string("/Users/bt/secrets/midjourney/user_id.txt")?;
     let user_id = MidjourneyUserId::from_string(user_id);
 
-    let result = imagine(ImagineRequest {
-      cookie_header,
-      hostname: MidjourneyHostname::Standard,
-      user_id: &user_id,
-      page_size: None,
+    let result = imagine(ImagineArgs {
+      request: ImagineRequest {
+        user_id: &user_id,
+        page_size: None,
+      },
+      cookie_header: &cookie_header,
+      hostname: None,
+      browser: None,
     }).await?;
 
     println!("Response: {:?}\n\n", result);

crates/api_clients/midjourney_client/src/endpoints/mod.rs

@@ -1,5 +1,7 @@
 pub mod get_index_page_html;
 pub mod imagine;
+pub mod prompt_session_log;
+pub mod refresh_token;
 pub mod storage_list;
 pub mod submit_job;
-pub mod user_connections;
+pub mod user_connections;

crates/api_clients/midjourney_client/src/endpoints/prompt_session_log.rs

@@ -0,0 +1,108 @@
+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, Serialize};
+
+/// The semantic parameters of a prompt-session-log entry.
+///
+/// This endpoint is telemetry the web client fires alongside each imagine
+/// submission (see the `08`/`11` captures); it is NOT the job-enqueue call
+/// (that is `submit_job`). We mirror it so our traffic matches the browser's.
+pub struct PromptSessionLogRequest<'a> {
+  /// 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<BrowserProfile>,
+}
+
+#[derive(Debug, Clone)]
+pub struct PromptSessionLogResponse {
+  /// The `status` field from the response body (e.g. `"success"`).
+  pub status: Option<String>,
+}
+
+pub async fn prompt_session_log(args: PromptSessionLogArgs<'_>) -> Result<PromptSessionLogResponse, MidjourneyError> {
+  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<String>,
+  }
+
+  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<String>,
+  }
+
+  let response = serde_json::from_str::<RawResponse>(&response_body).map_err(|err| MidjourneyApiError::deserialization(err, response_body.as_str()))?;
+
+  Ok(PromptSessionLogResponse { status: response.status })
+}

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<BrowserProfile>,
+}
+
+#[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<String>,
+
+  pub token_type: Option<String>,
+  pub user_id: Option<String>,
+  pub project_id: Option<String>,
+}
+
+pub async fn refresh_token(args: RefreshTokenArgs<'_>) -> Result<RefreshTokenResponse, MidjourneyError> {
+  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<String>,
+    token_type: Option<String>,
+    user_id: Option<String>,
+    project_id: Option<String>,
+  }
+
+  let response = serde_json::from_str::<RawResponse>(&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 })
+}

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<BrowserProfile>,
 }
 
 #[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<Vec<StorageItem>, MidjourneyError> {
+pub async fn storage_list(args: StorageListArgs<'_>) -> Result<Vec<StorageItem>, 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<Vec<StorageItem>
       .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<Vec<StorageItem>
   }
 
   let response : Vec<RawStorageItem> = 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<Vec<StorageItem>
 
 #[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);

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<BrowserProfile>,
 }
 
 #[derive(Debug, Clone)]
 pub struct SubmitJobResponse {
   /// On success, the job ID is returned.
   pub maybe_job_id: Option<String>,
-  
+
   /// On error, we have a list of error messages.
   pub maybe_errors: Option<Vec<SubmitJobError>>,
 }
 
+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<String>,
+  pub error_type: MidjourneySubmitErrorType,
   pub message: Option<String>,
 }
 
-pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse, MidjourneyError> {
-  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<SubmitJobResponse, MidjourneyError> {
+  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<SubmitJobResponse,
       .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")
@@ -72,11 +121,13 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse,
   #[derive(Serialize)]
   #[allow(non_snake_case)]
   struct Metadata {
+    // NB: Field order and null-vs-empty match the browser request exactly.
+    isMobile: Option<bool>,
     imagePrompts: u8,
     imageReferences: u8,
     characterReferences: u8,
     depthReferences: u8,
-    lightboxOpen: String,
+    lightboxOpen: Option<bool>,
   }
 
   #[derive(Serialize)]
@@ -94,16 +145,17 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse,
       mode: "fast".to_string(),
       private: false,
     },
-    channelId: req.channel_id.to_string(),
+    channelId: args.request.channel_id.to_string(),
     metadata: Metadata {
+      isMobile: None,
       imagePrompts: 0,
       imageReferences: 0,
       characterReferences: 0,
       depthReferences: 0,
-      lightboxOpen: "".to_string(),
+      lightboxOpen: None,
     },
     t: "imagine".to_string(),
-    prompt: req.prompt.to_string(),
+    prompt: args.request.prompt.to_string(),
   };
 
   let http_request  = http_request
@@ -126,6 +178,13 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse,
   let response_body = &response.text().await
       .map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
 
+  log::info!(
+    "Midjourney submit-jobs response: status={}, body_len={}, looks_like_challenge={}",
+    status.as_u16(),
+    response_body.len(),
+    response_body.contains("challenge-platform") || response_body.contains("Just a moment"),
+  );
+
   if !status.is_success() {
     if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
       return Err(MidjourneyApiError::CloudflareError(err).into());
@@ -170,7 +229,7 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse,
   }
 
   let response = serde_json::from_str::<RawResponse>(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<SubmitJobResponse,
   if let Some(failures) = response.failure.as_ref() {
     if !failures.is_empty() {
       maybe_errors = Some(failures.iter()
-          .map(|f| SubmitJobError { 
-            error_type: Some(f.r#type.clone()), 
-            message: Some(f.message.clone()), 
+          .map(|f| SubmitJobError {
+            error_type: MidjourneySubmitErrorType::from_api_str(&f.r#type),
+            message: Some(f.message.clone()),
           }).collect());
     }
   }
@@ -200,11 +259,40 @@ pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse,
 
 #[cfg(test)]
 mod tests {
-  use crate::client::midjourney_hostname::MidjourneyHostname;
-  use crate::endpoints::submit_job::{submit_job, SubmitJobRequest};
+  use crate::endpoints::submit_job::{
+    submit_job, MidjourneySubmitErrorType, SubmitJobArgs, SubmitJobError, SubmitJobRequest,
+    SubmitJobResponse,
+  };
   use errors::AnyhowResult;
   use filesys::read_to_trimmed_string::read_to_trimmed_string;
 
+  #[test]
+  fn classifies_subscription_required() {
+    assert_eq!(
+      MidjourneySubmitErrorType::from_api_str("subscription_required"),
+      MidjourneySubmitErrorType::SubscriptionRequired,
+    );
+    assert_eq!(
+      MidjourneySubmitErrorType::from_api_str("softban"),
+      MidjourneySubmitErrorType::Other("softban".to_string()),
+    );
+  }
+
+  #[test]
+  fn response_flags_subscription_required() {
+    let response = SubmitJobResponse {
+      maybe_job_id: None,
+      maybe_errors: Some(vec![SubmitJobError {
+        error_type: MidjourneySubmitErrorType::SubscriptionRequired,
+        message: Some("Please subscribe".to_string()),
+      }]),
+    };
+    assert!(response.is_subscription_required());
+
+    let ok = SubmitJobResponse { maybe_job_id: Some("job".to_string()), maybe_errors: None };
+    assert!(!ok.is_subscription_required());
+  }
+
   // Get channel id via:
   // https://identitytoolkit.googleapis.com/v1/accounts:lookup?key=[TOKEN]
   #[ignore]
@@ -213,11 +301,14 @@ mod tests {
     let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
     let channel_id = read_to_trimmed_string("/Users/bt/secrets/midjourney/channel_id.txt")?;
 
-    let result = submit_job(SubmitJobRequest {
-      prompt: "a modern n64 console",
-      channel_id: &channel_id,
-      cookie_header,
-      hostname: MidjourneyHostname::Standard,
+    let result = submit_job(SubmitJobArgs {
+      request: SubmitJobRequest {
+        prompt: "a modern n64 console",
+        channel_id: &channel_id,
+      },
+      cookie_header: &cookie_header,
+      hostname: None,
+      browser: None,
     }).await?;
 
     println!("Response: {:?}", result);

crates/api_clients/midjourney_client/src/endpoints/user_connections.rs

@@ -2,17 +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;
 
-const USER_CONNECTIONS_URL :&str = "https://www.midjourney.com/api/user-connections";
-
-pub struct GetUserConnectionsRequest {
-  pub hostname: MidjourneyHostname,
-  pub cookie_header: String,
+/// Lists the user's linked auth providers. Has no semantic parameters — only
+/// transport concerns.
+pub struct UserConnectionsArgs<'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<BrowserProfile>,
 }
 
 #[derive(Debug)]
@@ -30,29 +32,28 @@ pub struct Provider {
   pub is_linked: bool,
 }
 
-pub async fn user_connections(req: GetUserConnectionsRequest) -> Result<Vec<Provider>, MidjourneyError> {
+pub async fn user_connections(args: UserConnectionsArgs<'_>) -> Result<Vec<Provider>, 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<Vec<Prov
       .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")
@@ -128,7 +128,7 @@ pub async fn user_connections(req: GetUserConnectionsRequest) -> Result<Vec<Prov
   }
 
   let response : Vec<RawProvider> = 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<Vec<Prov
 
 #[cfg(test)]
 mod tests {
-  use crate::client::midjourney_hostname::MidjourneyHostname;
-  use crate::endpoints::user_connections::{user_connections, GetUserConnectionsRequest};
+  use crate::endpoints::user_connections::{user_connections, UserConnectionsArgs};
   use errors::AnyhowResult;
   use filesys::read_to_trimmed_string::read_to_trimmed_string;
 
@@ -159,9 +158,10 @@ mod tests {
   async fn test() -> 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);

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<serde_json::Error> 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: "<redacted>".to_string() }
   }
 }
 

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."),
     }
   }
 }

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<BrowserProfile>,
 }
 
-pub async fn get_midjourney_user_id_from_storage(req: GetMidjourneyUserIdRequest) -> Result<MidjourneyUserId, MidjourneyError> {
-  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<MidjourneyUserId, MidjourneyError> {
+  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(())
   }
 }
-

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<BrowserProfile>,
 }
 
 #[derive(Debug, Clone)]
@@ -18,12 +24,25 @@ pub struct GetUserInfoResponse {
   pub websocket_token: Option<String>,
 }
 
-pub async fn get_user_info(req: GetUserInfoRequest) -> Result<GetUserInfoResponse, MidjourneyError> {
-  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<GetUserInfoResponse, MidjourneyError> {
+  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"),
+  );
+
   /*
     <script id="initialProps" type="application/json">
       {
@@ -107,7 +126,7 @@ pub async fn get_user_info(req: GetUserInfoRequest) -> Result<GetUserInfoRespons
   }
 
   let response : RawBody = serde_json::from_str(&inner_json)
-      .map_err(|err| MidjourneyApiError::DeserializationError(err))?;
+      .map_err(|err| MidjourneyApiError::deserialization(err, inner_json.as_str()))?;
   
   let props = match response.initialAuthUser {
     Some(props) => props,
@@ -128,8 +147,7 @@ pub async fn get_user_info(req: GetUserInfoRequest) -> Result<GetUserInfoRespons
 
 #[cfg(test)]
 mod tests {
-  use crate::client::midjourney_hostname::MidjourneyHostname;
-  use crate::recipes::get_user_info::{get_user_info, GetUserInfoRequest};
+  use crate::recipes::get_user_info::{get_user_info, GetUserInfoArgs};
   use errors::AnyhowResult;
   use filesys::read_to_trimmed_string::read_to_trimmed_string;
 
@@ -138,9 +156,10 @@ mod tests {
   async fn test() -> AnyhowResult<()> {
     let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
 
-    let result = get_user_info(GetUserInfoRequest {
-      cookie_header,
-      hostname: MidjourneyHostname::Standard,
+    let result = get_user_info(GetUserInfoArgs {
+      cookie_header: &cookie_header,
+      hostname: None,
+      browser: None,
     }).await?;
 
     println!("Response: {:?}\n\n", result);

crates/api_clients/midjourney_client/src/recipes/text_to_image.rs

@@ -1,13 +1,23 @@
 use crate::client::midjourney_hostname::MidjourneyHostname;
-use crate::endpoints::submit_job::{submit_job, SubmitJobRequest};
+use crate::endpoints::submit_job::{submit_job, MidjourneySubmitErrorType, SubmitJobArgs, SubmitJobRequest};
 use crate::error::midjourney_error::MidjourneyError;
 use crate::recipes::channel_id::ChannelId;
+use browser_emulation::browser_profile::BrowserProfile;
 
+/// The semantic parameters of a text-to-image request.
 pub struct TextToImageRequest<'a> {
   pub prompt: &'a str,
   pub channel_id: &'a ChannelId,
-  pub hostname: MidjourneyHostname,
-  pub cookie_header: String,
+}
+
+/// A text-to-image request plus its transport concerns.
+pub struct TextToImageArgs<'a> {
+  pub request: TextToImageRequest<'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<BrowserProfile>,
 }
 
 #[derive(Debug, Clone)]
@@ -19,22 +29,39 @@ pub struct TextToImageResponse {
   pub maybe_errors: Option<Vec<TextToImageError>>,
 }
 
+impl TextToImageResponse {
+  pub fn errors(&self) -> &[TextToImageError] {
+    self.maybe_errors.as_deref().unwrap_or(&[])
+  }
+
+  /// Whether the request was rejected for lack of an active paid Midjourney
+  /// subscription. See [`MidjourneySubmitErrorType::SubscriptionRequired`].
+  pub fn is_subscription_required(&self) -> bool {
+    self.errors()
+        .iter()
+        .any(|error| error.error_type == MidjourneySubmitErrorType::SubscriptionRequired)
+  }
+}
+
 #[derive(Debug, Clone)]
 pub struct TextToImageError {
-  pub error_type: Option<String>,
+  pub error_type: MidjourneySubmitErrorType,
   pub message: Option<String>,
 }
 
 /// Slightly more ergonomic text-to-image API.
 /// As we add more `submit_job()` cases, we'll keep this simple.
-pub async fn text_to_image(req: TextToImageRequest<'_>) -> Result<TextToImageResponse, MidjourneyError> {
-  let channel_id = req.channel_id.to_string();
+pub async fn text_to_image(args: TextToImageArgs<'_>) -> Result<TextToImageResponse, MidjourneyError> {
+  let channel_id = args.request.channel_id.to_string();
 
-  let response = submit_job(SubmitJobRequest {
-    prompt: req.prompt,
-    channel_id: &channel_id,
-    hostname: req.hostname,
-    cookie_header: req.cookie_header,
+  let response = submit_job(SubmitJobArgs {
+    request: SubmitJobRequest {
+      prompt: args.request.prompt,
+      channel_id: &channel_id,
+    },
+    cookie_header: args.cookie_header,
+    hostname: args.hostname,
+    browser: args.browser,
   }).await?;
 
   Ok(TextToImageResponse {

crates/api_clients/midjourney_client/src/utils/image_downloader_client.rs

@@ -2,9 +2,9 @@ use crate::error::midjourney_api_error::MidjourneyApiError;
 use crate::error::midjourney_client_error::MidjourneyClientError;
 use crate::error::midjourney_error::MidjourneyError;
 use crate::utils::get_image_url::get_image_url;
+use browser_emulation::browser_profile::BrowserProfile;
 use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
 use wreq::Client;
-use wreq_util::Emulation;
 
 #[derive(Clone)]
 pub struct ImageDownloaderClient {
@@ -12,12 +12,12 @@ pub struct ImageDownloaderClient {
 }
 
 impl ImageDownloaderClient {
-  pub fn create() -> Result<Self, MidjourneyClientError> {
+  /// Build a downloader. `maybe_browser` defaults to [`BrowserProfile::default`].
+  pub fn create(maybe_browser: Option<BrowserProfile>) -> Result<Self, MidjourneyClientError> {
     Ok(Self {
-      client: Client::builder()
-          .emulation(Emulation::Firefox139)
-          .build()
-          .map_err(|err| MidjourneyClientError::WreqError(err))?
+      client: maybe_browser.unwrap_or_default()
+          .build_client()
+          .map_err(MidjourneyClientError::WreqError)?,
     })
   }
 
@@ -26,16 +26,13 @@ impl ImageDownloaderClient {
 
     // TODO: Cookies
     // TODO: Cache control headers?
-    let mut http_request = self.client.get(url)
-        //.header("Referrer-Policy", "origin-when-cross-origin")
-        //.header("content-type", "application/json")
-        //.header("sec-fetch-user", "?1")
-        //.header("upgrade-insecure-requests", "1");
+    // NB: Browser-identity headers come from the emulation on the client; only
+    // request-context headers are set here.
+    let http_request = self.client.get(url)
         .header("Referrer", "https://www.midjourney.com/")
         .header("accept", "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8")
         .header("accept-language", "en-US,en;q=0.8")
         .header("priority", "i")
-        .header("sec-ch-ua-mobile", "?0")
         .header("sec-fetch-dest", "image")
         .header("sec-fetch-mode", "no-cors")
         .header("sec-fetch-site", "same-site");

crates/desktop/artcraft/Cargo.toml

@@ -22,6 +22,9 @@ crate-type = ["staticlib", "cdylib", "rlib"]
 tauri-build = { version = "2.6.2", features = [] }
 
 [dependencies]
+toml.workspace = true
+browser_emulation.workspace = true
+cookie_store_wrapper = { workspace = true, features = ["wreq"] }
 
 # Internal
 artcraft_api_defs.workspace = true

crates/desktop/artcraft/src/core/commands/app_preferences/get_app_preferences_command.rs

@@ -1,5 +1,6 @@
 use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager;
 use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory;
+use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename;
 use errors::AnyhowResult;
 use log::{error, info};
 use serde_derive::Serialize;
@@ -14,6 +15,8 @@ pub struct GetAppPreferencesResponse {
 pub struct AppPreferencesPayload {
   /// The downloads directory to use when a user downloads a file.
   pub preferred_download_directory: PreferredDownloadDirectory,
+  pub preferred_download_filename: PreferredDownloadFilename,
+  pub auto_download: bool,
   
   /// Play sounds on events.
   pub play_sounds: bool,
@@ -63,6 +66,8 @@ async fn get_prefs(app_prefs: &AppPreferencesManager) -> AnyhowResult<AppPrefere
   let prefs = app_prefs.get_clone()?;
   Ok(AppPreferencesPayload {
     preferred_download_directory: prefs.preferred_download_directory,
+    preferred_download_filename: prefs.preferred_download_filename,
+    auto_download: prefs.auto_download,
     play_sounds: prefs.play_sounds,
     delete_file_sound: prefs.delete_file_sound,
     enqueue_success_sound: prefs.enqueue_success_sound,

crates/desktop/artcraft/src/core/commands/app_preferences/update_app_preference_command.rs

@@ -1,5 +1,6 @@
 use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager;
 use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory;
+use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use anyhow::anyhow;
 use errors::AnyhowResult;
@@ -21,12 +22,15 @@ pub enum ValueType {
   Bool(bool),
   String(String),
   DownloadDirectory(PreferredDownloadDirectory),
+  DownloadFilename(PreferredDownloadFilename),
 }
 
 #[derive(Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum PreferenceName {
   PreferredDownloadDirectory,
+  PreferredDownloadFilename,
+  AutoDownload,
   PlaySounds,
   DeleteFileSound,
   EnqueueSuccessSound,
@@ -70,6 +74,15 @@ async fn update_prefs(
   info!("Value is: {:?}", request.value);
   
   match request.preference {
+    PreferenceName::AutoDownload => {
+      match request.value {
+        Some(ValueType::Bool(value)) => prefs.auto_download = value,
+        _ => return Err(anyhow!("Auto Download must be a boolean")),
+      }
+    }
+    PreferenceName::PreferredDownloadFilename => {
+      prefs.preferred_download_filename = filename_value(request.value)?;
+    }
     PreferenceName::PreferredDownloadDirectory => {
       match request.value {
         Some(ValueType::DownloadDirectory(dir)) => 
@@ -113,8 +126,8 @@ async fn update_prefs(
     }
   }
   
-  app_prefs.set_clone(&prefs)?;
   app_data_root.settings_dir().write_app_preferences(&prefs)?;
+  app_prefs.set_clone(&prefs)?;
   
   Ok(())
 }
@@ -125,3 +138,48 @@ fn string_value(value: &ValueType) -> AnyhowResult<String> {
     _ => Err(anyhow!("Invalid value type: {:?}", value)),
   }
 }
+
+fn filename_value(value: Option<ValueType>) -> AnyhowResult<PreferredDownloadFilename> {
+  // The untagged value enum decodes string variants as String first.
+  let filename = match value {
+    Some(ValueType::String(value)) if value == "artcraft_convention" => PreferredDownloadFilename::ArtcraftConvention,
+    Some(ValueType::DownloadFilename(filename)) => filename,
+    _ => return Err(anyhow!("Invalid download filename preference")),
+  };
+  if let PreferredDownloadFilename::Custom(format) = &filename {
+    PreferredDownloadFilename::validate_custom_format(format).map_err(|err| anyhow!(err))?;
+  }
+  Ok(filename)
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use crate::core::state::app_preferences::app_preferences::AppPreferences;
+
+  #[tokio::test]
+  async fn download_preferences_accept_wire_values_and_survive_reloading() {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    let manager = AppPreferencesManager::with_owned(AppPreferences::default());
+    for request in [
+      serde_json::json!({ "preference": "preferred_download_directory", "value": {"custom": "/existing/downloads"} }),
+      serde_json::json!({ "preference": "auto_download", "value": true }),
+      serde_json::json!({ "preference": "preferred_download_filename", "value": {"custom_format": "my_{model}_{date}"} }),
+    ] {
+      update_prefs(serde_json::from_value(request).unwrap(), &manager, &root).await.unwrap();
+    }
+    let prefs = AppPreferences::load_from_file_or_default(&root);
+    assert!(prefs.auto_download);
+    assert_eq!(prefs.preferred_download_directory, PreferredDownloadDirectory::Custom("/existing/downloads".into()));
+    assert_eq!(prefs.preferred_download_filename, PreferredDownloadFilename::Custom("my_{model}_{date}".into()));
+
+    let invalid = serde_json::json!({ "preference": "preferred_download_filename", "value": {"custom_format": "../escape"} });
+    assert!(update_prefs(serde_json::from_value(invalid).unwrap(), &manager, &root).await.is_err());
+    assert_eq!(AppPreferences::load_from_file_or_default(&root).preferred_download_filename, prefs.preferred_download_filename);
+    let reset = serde_json::json!({ "preference": "preferred_download_filename", "value": "artcraft_convention" });
+    update_prefs(serde_json::from_value(reset).unwrap(), &manager, &root).await.unwrap();
+    assert_eq!(AppPreferences::load_from_file_or_default(&root).preferred_download_filename, PreferredDownloadFilename::ArtcraftConvention);
+    assert!(manager.get_clone().unwrap().auto_download);
+  }
+}

crates/desktop/artcraft/src/core/commands/cost_estimate/estimate_image_cost_command.rs

@@ -1,3 +1,5 @@
+use enums::common::generation_provider::GenerationProvider;
+use crate::core::commands::generate::omni::OmniResponse;
 use crate::core::commands::generate::omni::{self, Modality, OmniRequest, OmniResult};
 use crate::core::commands::generate::omni::dispatch::{adapt_legacy_response, decode_native};
 use tauri::{AppHandle, Manager};
@@ -17,6 +19,23 @@ impl SerializeMarker for EstimateImageCostResponse {}
 
 #[tauri::command]
 pub async fn estimate_image_cost_command(request: OmniRequest, app: AppHandle) -> OmniResult {
+  if matches!(request.provider, Some(GenerationProvider::Midjourney)) {
+    if !matches!(request.model(), Some("midjourney" | "midjourney_7" | "midjourney_7_niji" | "midjourney_8")) {
+      return Err(CommandErrorResponseWrapper {
+        status: CommandErrorStatus::BadRequest,
+        error_message: Some("This model cannot use a Midjourney account".to_string()),
+        error_type: Some("invalid_provider_for_model".to_string()),
+        error_details: None,
+      });
+    }
+    // Paid for by the user's Midjourney subscription, with no ArtCraft credit
+    // charge. Do not send this estimate to ArtCraft's generation backend.
+    return Ok(OmniResponse(serde_json::json!({
+      "success": true, "cost_in_credits": 0, "cost_in_usd_cents": null,
+      "is_free": false, "is_unlimited": false, "is_rate_limited": true,
+      "has_watermark": false,
+    })).into());
+  }
   if request.uses_artcraft() && !request.uses_legacy_image_endpoint() {
     return omni::estimate(request, Modality::Image, &app).await;
   }

crates/desktop/artcraft/src/core/commands/deprecated/text_to_image/midjourney/handle_midjourney.rs

@@ -1,143 +1,18 @@
-use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason, ProviderFailureReason};
-use crate::core::commands::deprecated::image_edit::enqueue_edit_image_command::{EditImageQuality, EditImageSize};
+use crate::core::commands::deprecated::text_to_image::enqueue_text_to_image_command::EnqueueTextToImageRequest;
+use crate::core::commands::enqueue::generate_error::GenerateError;
 use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
-use crate::core::commands::deprecated::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize};
-use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
-use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent;
-use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent;
-use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
-use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
-use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent;
+use crate::core::commands::generate::generate_image::providers::midjourney::handle_midjourney::handle_midjourney as handle_midjourney_generation;
+use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
-use crate::core::state::data_dir::app_data_root::AppDataRoot;
-use crate::core::state::provider_priority::ProviderPriorityStore;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
-use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
-use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
-use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
-use enums::common::generation_provider::GenerationProvider;
-use enums::tauri::tasks::task_type::TaskType;
-use uuid_utils::uuid::generate_random_uuid;
-use log::{error, info};
-use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
-use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest};
-use midjourney_client::error::midjourney_api_error::MidjourneyApiError;
-use midjourney_client::recipes::channel_id::ChannelId;
-use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageError, TextToImageRequest};
-use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
-use openai_sora_client::recipes::simple_image_gen_with_session_auto_renew::{simple_image_gen_with_session_auto_renew, SimpleImageGenAutoRenewRequest};
-use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages};
-use std::time::Duration;
-use tauri::AppHandle;
-use tokens::tokens::media_files::MediaFileToken;
+use tauri::{AppHandle, Manager};
 
-pub async fn handle_midjourney(
-  app: &AppHandle,
-  request: &EnqueueTextToImageRequest,
-  app_env_configs: &AppEnvConfigs,
-  mj_creds_manager: &MidjourneyCredentialManager,
-) -> Result<TaskEnqueueSuccess, GenerateError> {
-
-  let creds = match mj_creds_manager.maybe_copy_cookie_store() {
-    Ok(Some(creds)) => creds,
-    Ok(None) => {
-      error!("Midjourney credentials not found.");
-      ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Midjourney, &app);
-      return Err(GenerateError::needs_midjourney_credentials());
-    }
-    Err(err) => {
-      error!("Error reading Midjourney credentials: {:?}", err);
-      ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Midjourney, &app);
-      return Err(GenerateError::needs_midjourney_credentials());
-    },
-  };
-
-  // TODO: We can request population of the user info if absent or expired.
-  
-  let user_info = match mj_creds_manager.maybe_copy_user_info() {
-    Ok(Some(user_info)) => user_info,
-    Ok(None) => {
-      return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyUserInfo));
-    }
-    Err(err) => {
-      error!("Error reading Midjourney user info: {:?}", err);
-      return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyUserInfo));
-    },
-  };
-
-  let channel_id = match user_info.user_id {
-    Some(user_id) => ChannelId::UserId(user_id),
-    None => {
-      error!("Midjourney user info does not contain a user ID.");
-      return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyUserId));
-    }
-  };
-
-  info!("Calling midjourney ...");
-
-  let cookie_header = creds.to_cookie_string();
-
-  let prompt = request.prompt
-      .as_deref()
-      .unwrap_or("");
-
-  let result = text_to_image(TextToImageRequest {
-    prompt,
-    channel_id: &channel_id,
-    hostname: MidjourneyHostname::Standard,
-    cookie_header,
-  }).await;
-
-  let result = match result {
-    Ok(result) => result,
-    Err(err) => {
-      error!("Failed to use MidJourney: {:?}", err);
-      return Err(GenerateError::from(err));
-    }
-  };
-  
-  let job_id = match result.maybe_job_id {
-    Some(job_id) => job_id,
-    None => {
-      error!("Failed to enqueue MidJourney: No job ID returned.");
-      return handle_midjourney_errors(app, result.maybe_errors);
-    }
-  };
-
-  info!("Successfully enqueued MidJourney. Job token: {}", job_id);
-
-  Ok(TaskEnqueueSuccess {
-    provider: GenerationProvider::Midjourney,
-    model: Some(GenerationModel::Midjourney),
-    provider_job_id: Some(job_id),
-    task_type: TaskType::ImageGeneration,
-    maybe_queue_status_url: None,
-    maybe_prompt_token: None,
-    maybe_queue_response_url: None,
-  })
-}
-
-fn handle_midjourney_errors(
-  app: &AppHandle,
-  maybe_errors: Option<Vec<TextToImageError>>
-) -> Result<TaskEnqueueSuccess, GenerateError> {
-  if let Some(errors) = maybe_errors {
-    if !errors.is_empty() {
-      let messages: Vec<String> = errors.iter()
-          .map(|e| format!("{:?}", e))
-          .collect();
-
-      let combined_message = messages.join("; ");
-
-      let event = FlashUserInputErrorEvent {
-        message: format!("Midjourney Error: {}", combined_message),
-      };
-
-      if let Err(err) = event.send(&app) {
-        error!("Failed to send FlashUserInputErrorEvent: {:?}", err); // Fail open
-      }
-    }
-  }
-  
-  Err(GenerateError::ProviderFailure(ProviderFailureReason::MidjourneyJobEnqueueFailed))
+pub async fn handle_midjourney(app: &AppHandle, request: &EnqueueTextToImageRequest, config: &AppEnvConfigs, credentials: &MidjourneyCredentialManager) -> Result<TaskEnqueueSuccess, GenerateError> {
+  let request: TauriGenerateImageRequest = serde_json::from_value(serde_json::json!({
+    "model": "midjourney", "provider": "midjourney", "prompt": request.prompt,
+    "aspect_ratio": request.common_aspect_ratio, "batch_size": request.number_images,
+    "image_media_tokens": request.image_media_tokens,
+  }))
+  .map_err(|err| GenerateError::AnyhowError(err.into()))?;
+  handle_midjourney_generation(&request, config, credentials, &app.state()).await
 }

crates/desktop/artcraft/src/core/commands/download/download_url_command.rs

@@ -10,7 +10,8 @@ use crate::core::state::app_preferences::app_preferences::AppPreferences;
 use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::utils::download_url_to_temp_dir::download_url_to_temp_dir;
-use crate::core::utils::download_url_to_user_download_dir::download_url_to_user_download_dir;
+use crate::core::utils::download_url_to_user_download_dir::{download_url_with_filename, suggested_download_filename};
+use chrono::Local;
 use anyhow::anyhow;
 use log::{error, info};
 use serde_derive::{Deserialize, Serialize};
@@ -23,6 +24,13 @@ use url::Url;
 #[derive(Deserialize, Debug)]
 pub struct DownloadUrlRequest {
   pub url: Url,
+  #[serde(rename = "model")]
+  pub maybe_model: Option<String>,
+  #[serde(rename = "batch_index")]
+  pub maybe_batch_index: Option<usize>,
+  /// Set only when the user explicitly chooses a folder for a batch download.
+  #[serde(rename = "directory")]
+  pub maybe_directory: Option<PathBuf>,
 }
 
 #[derive(Serialize)]
@@ -113,11 +121,9 @@ pub async fn handle_request(
 
   let app_prefs = app_prefs.get_clone()?;
 
-  let download_path = download_url_to_user_download_dir(
-    &request.url,
-    app_data_root,
-    &app_prefs
-  ).await?;
+  let filename = suggested_download_filename(&request.url, &app_prefs, request.maybe_model.as_deref(), request.maybe_batch_index, Local::now());
+  let directory = request.maybe_directory.unwrap_or_else(|| app_prefs.preferred_download_directory.download_directory(app_data_root));
+  let download_path = download_url_with_filename(&request.url, &directory, &filename).await?;
 
   info!("downloaded to: {:?}", download_path);
 

crates/desktop/artcraft/src/core/commands/download/get_download_path_command.rs

@@ -0,0 +1,46 @@
+use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager;
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::utils::download_url_to_user_download_dir::suggested_download_filename;
+use chrono::Local;
+use serde::{Deserialize, Serialize};
+use std::path::PathBuf;
+use tauri::State;
+use url::Url;
+
+#[derive(Deserialize)]
+pub struct GetDownloadPathRequest {
+  pub url: Url,
+  #[serde(rename = "model")]
+  pub maybe_model: Option<String>,
+  #[serde(rename = "batch_index")]
+  pub maybe_batch_index: Option<usize>,
+}
+
+#[derive(Serialize)]
+pub struct GetDownloadPathResponse {
+  pub filename: String,
+  pub path: PathBuf,
+}
+
+/// Suggest a name for a native save dialog using the same formatter as
+/// automatic downloads. The dialog may still let the user override it.
+#[tauri::command]
+pub fn get_download_path_command(
+  request: GetDownloadPathRequest,
+  app_prefs: State<'_, AppPreferencesManager>,
+  app_data_root: State<'_, AppDataRoot>,
+) -> Result<GetDownloadPathResponse, String> {
+  let prefs = app_prefs.get_clone().map_err(|err| err.to_string())?;
+  let filename = suggested_download_filename(
+    &request.url,
+    &prefs,
+    request.maybe_model.as_deref(),
+    request.maybe_batch_index,
+    Local::now(),
+  );
+  let path = prefs
+    .preferred_download_directory
+    .download_directory(&app_data_root)
+    .join(&filename);
+  Ok(GetDownloadPathResponse { filename, path })
+}

crates/desktop/artcraft/src/core/commands/download/mod.rs

@@ -1,3 +1,4 @@
 pub mod download_directory_reveal_command;
 pub mod download_media_file_command;
 pub mod download_url_command;
+pub mod get_download_path_command;

crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs

@@ -233,6 +233,7 @@ impl From<ArtcraftRouterError> for GenerateError {
       ArtcraftRouterError::InvalidInput(msg) => Self::BadInput(BadInputReason::WrongImageArguments(msg)),
       ArtcraftRouterError::ProviderBillingError(err) => {
         let provider = match err {
+          ProviderError::Midjourney(_) | ProviderError::MidjourneySubscriptionRequired(_) | ProviderError::MidjourneySubmitRejected(_) => BillingProvider::Midjourney,
           ProviderError::Fal(_) => BillingProvider::Fal,
           ProviderError::KinoviWeb(_) => BillingProvider::Kinovi,
           ProviderError::Storyteller(_) => BillingProvider::Artcraft,
@@ -242,6 +243,8 @@ impl From<ArtcraftRouterError> for GenerateError {
         };
         Self::BillingIssue(BillingIssueReason { provider })
       },
+      ArtcraftRouterError::Provider(ProviderError::Midjourney(e)) => Self::from(e),
+      ArtcraftRouterError::Provider(ProviderError::MidjourneySubscriptionRequired(message) | ProviderError::MidjourneySubmitRejected(message)) => Self::BadInput(BadInputReason::WrongImageArguments(message)),
       ArtcraftRouterError::Provider(ProviderError::Storyteller(e)) => Self::ProviderFailure(ProviderFailureReason::StorytellerError(e)),
       ArtcraftRouterError::Provider(ProviderError::Fal(_)) => Self::FalNoLongerSupported,
       ArtcraftRouterError::Provider(ProviderError::GmiCloud(_)) => Self::ArtcraftRouterNotYetSupportedProvider("gmicloud"),

crates/desktop/artcraft/src/core/commands/generate/common/router_image_request_to_artcraft_prompt.rs

@@ -71,6 +71,7 @@ fn image_model_to_common_model_type(model: RouterImageModel) -> Option<CommonMod
 fn provider_to_generation_provider(provider: RouterProvider) -> GenerationProvider {
   match provider {
     RouterProvider::Artcraft => GenerationProvider::Artcraft,
+    RouterProvider::Midjourney => GenerationProvider::Midjourney,
     RouterProvider::Fal => GenerationProvider::Fal,
     // Unused providers -> ArtCraft
     RouterProvider::KinoviWeb => GenerationProvider::Artcraft ,

crates/desktop/artcraft/src/core/commands/generate/common/router_video_request_to_artcraft_prompt.rs

@@ -96,6 +96,7 @@ fn video_model_to_common_model_type(model: RouterVideoModel) -> Option<CommonMod
 fn provider_to_generation_provider(provider: RouterProvider) -> GenerationProvider {
   match provider {
     RouterProvider::Artcraft => GenerationProvider::Artcraft,
+    RouterProvider::Midjourney => GenerationProvider::Midjourney,
     RouterProvider::Fal => GenerationProvider::Fal,
     // Unused providers -> ArtCraft
     RouterProvider::KinoviWeb => GenerationProvider::Artcraft ,

crates/desktop/artcraft/src/core/commands/generate/generate_image/generate_image_command.rs

@@ -1,3 +1,5 @@
+use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent;
+use crate::core::commands::generate::generate_image::providers::midjourney::handle_midjourney::handle_midjourney;
 use crate::core::commands::generate::omni::{self, Modality, OmniRequest, OmniResult};
 use crate::core::commands::generate::omni::dispatch::{adapt_legacy_response, decode_native};
 use tauri::Manager;
@@ -54,11 +56,8 @@ async fn generate_image_native(
         &storyteller_creds_manager,
       ).await
     }
-    // Midjourney uses its own legacy command path, not this one.
     GenerationProvider::Midjourney => {
-      Err(GenerateError::NotYetImplemented(
-        "Midjourney should use its dedicated command".to_string(),
-      ))
+      handle_midjourney(&request, &app_env_configs, &app.state(), &storyteller_creds_manager).await
     }
     // All other providers go through the router.
     other => {
@@ -137,11 +136,15 @@ async fn handle_error_behavior(
   error!("generate_image_command error: {:?}", err);
 
   notify_frontend_of_errors(&app, &err).await;
+  if matches!(&err, GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyCredentials | MissingCredentialsReason::NeedsMidjourneyUserId | MissingCredentialsReason::NeedsMidjourneyUserInfo)) {
+    ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Midjourney, app);
+  }
 
   let error_type = match &err {
     GenerateError::BadInput(_) => TauriGenerateImageErrorType::BadInput,
     GenerateError::MissingCredentials(MissingCredentialsReason::NeedsFalApiKey) => TauriGenerateImageErrorType::NeedsFalApiKey,
     GenerateError::MissingCredentials(MissingCredentialsReason::NeedsGrokCredentials) => TauriGenerateImageErrorType::NeedsGrokCredentials,
+    GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyCredentials | MissingCredentialsReason::NeedsMidjourneyUserId | MissingCredentialsReason::NeedsMidjourneyUserInfo) => TauriGenerateImageErrorType::NeedsMidjourneyCredentials,
     GenerateError::MissingCredentials(_) => TauriGenerateImageErrorType::NeedsStorytellerCredentials,
     GenerateError::NoProviderAvailable => TauriGenerateImageErrorType::NoProviderAvailable,
     GenerateError::BillingIssue(_) => TauriGenerateImageErrorType::BillingIssue,

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/artcraft_router/mod.rs

@@ -1,4 +1,4 @@
 pub mod handle_router;
 mod handle_api_providers;
 mod handle_web_login_providers;
-mod utils;
+pub(super) mod utils;

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/midjourney/handle_midjourney.rs

@@ -1,140 +1,60 @@
-use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason, ProviderFailureReason};
-use crate::core::commands::deprecated::image_edit::enqueue_edit_image_command::{EditImageQuality, EditImageSize};
+use crate::core::api_adapters::models::image::tauri_image_model_to_generation_model::tauri_image_model_to_generation_model;
+use crate::core::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model;
+use crate::core::commands::enqueue::generate_error::{BadInputReason, GenerateError};
 use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
-use crate::core::commands::deprecated::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize};
-use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
-use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent;
-use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent;
-use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
-use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
-use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent;
+use crate::core::commands::generate::common::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
+use crate::core::commands::generate::generate_image::providers::artcraft_router::utils::convert_enums_to_router::convert_aspect_ratio;
+use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
+use crate::core::commands::generate::generate_image::tauri_image_model::TauriImageModel;
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
-use crate::core::state::data_dir::app_data_root::AppDataRoot;
-use crate::core::state::provider_priority::ProviderPriorityStore;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
-use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
-use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use artcraft_client::endpoints::prompts::create_prompt::create_prompt;
+use artcraft_router::api::router_image_model::RouterImageModel;
+use artcraft_router::api::router_provider::RouterProvider;
+use artcraft_router::client::generation_mode_mismatch_strategy::GenerationModeMismatchStrategy;
+use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
+use artcraft_router::client::router_client::RouterClient;
+use artcraft_router::client::router_midjourney_client::RouterMidjourneyClient;
+use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder;
+use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse;
+use artcraft_router::generate::generate_image::image_generation_draft_or_request::ImageGenerationDraftOrRequest;
 use enums::common::generation_provider::GenerationProvider;
 use enums::tauri::tasks::task_type::TaskType;
-use uuid_utils::uuid::generate_random_uuid;
-use log::{error, info};
-use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
-use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest};
-use midjourney_client::error::midjourney_api_error::MidjourneyApiError;
-use midjourney_client::recipes::channel_id::ChannelId;
-use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageError, TextToImageRequest};
-use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
-use openai_sora_client::recipes::simple_image_gen_with_session_auto_renew::{simple_image_gen_with_session_auto_renew, SimpleImageGenAutoRenewRequest};
-use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages};
-use std::time::Duration;
-use tauri::AppHandle;
-use tokens::tokens::media_files::MediaFileToken;
 
-pub async fn handle_midjourney(
-  app: &AppHandle,
-  request: &EnqueueTextToImageRequest,
-  app_env_configs: &AppEnvConfigs,
-  mj_creds_manager: &MidjourneyCredentialManager,
-) -> Result<TaskEnqueueSuccess, GenerateError> {
-
-  let creds = match mj_creds_manager.maybe_copy_cookie_store() {
-    Ok(Some(creds)) => creds,
-    Ok(None) => {
-      error!("Midjourney credentials not found.");
-      ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Midjourney, &app);
-      return Err(GenerateError::needs_midjourney_credentials());
-    }
-    Err(err) => {
-      error!("Error reading Midjourney credentials: {:?}", err);
-      ShowProviderLoginModalEvent::send_for_provider(GenerationProvider::Midjourney, &app);
-      return Err(GenerateError::needs_midjourney_credentials());
-    },
-  };
-
-  // TODO: We can request population of the user info if absent or expired.
-  
-  let user_info = match mj_creds_manager.maybe_copy_user_info() {
-    Ok(Some(user_info)) => user_info,
-    Ok(None) => {
-      return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyUserInfo));
-    }
-    Err(err) => {
-      error!("Error reading Midjourney user info: {:?}", err);
-      return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyUserInfo));
-    },
-  };
-
-  let channel_id = match user_info.user_id {
-    Some(user_id) => ChannelId::UserId(user_id),
-    None => {
-      error!("Midjourney user info does not contain a user ID.");
-      return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsMidjourneyUserId));
-    }
-  };
-
-  info!("Calling midjourney ...");
-
-  let cookie_header = creds.to_cookie_string();
-
-  let prompt = request.prompt
-      .as_deref()
-      .unwrap_or("");
-
-  let result = text_to_image(TextToImageRequest {
-    prompt,
-    channel_id: &channel_id,
-    hostname: MidjourneyHostname::Standard,
-    cookie_header,
-  }).await;
-
-  let result = match result {
-    Ok(result) => result,
-    Err(err) => {
-      error!("Failed to use MidJourney: {:?}", err);
-      return Err(GenerateError::from(err));
-    }
-  };
-  
-  let job_id = match result.maybe_job_id {
-    Some(job_id) => job_id,
-    None => {
-      error!("Failed to enqueue MidJourney: No job ID returned.");
-      return handle_midjourney_errors(app, result.maybe_errors);
-    }
-  };
-
-  info!("Successfully enqueued MidJourney. Job token: {}", job_id);
-
-  Ok(TaskEnqueueSuccess {
-    provider: GenerationProvider::Midjourney,
-    model: Some(GenerationModel::Midjourney),
-    provider_job_id: Some(job_id),
-    task_type: TaskType::ImageGeneration,
-  })
-}
-
-fn handle_midjourney_errors(
-  app: &AppHandle,
-  maybe_errors: Option<Vec<TextToImageError>>
-) -> Result<TaskEnqueueSuccess, GenerateError> {
-  if let Some(errors) = maybe_errors {
-    if !errors.is_empty() {
-      let messages: Vec<String> = errors.iter()
-          .map(|e| format!("{:?}", e))
-          .collect();
-
-      let combined_message = messages.join("; ");
-
-      let event = FlashUserInputErrorEvent {
-        message: format!("Midjourney Error: {}", combined_message),
-      };
-
-      if let Err(err) = event.send(&app) {
-        error!("Failed to send FlashUserInputErrorEvent: {:?}", err); // Fail open
-      }
-    }
+pub async fn handle_midjourney(request: &TauriGenerateImageRequest, config: &AppEnvConfigs, credentials: &MidjourneyCredentialManager, storyteller: &StorytellerCredentialManager) -> Result<TaskEnqueueSuccess, GenerateError> {
+  let model = request.model.ok_or_else(GenerateError::no_model_specified)?;
+  if !matches!(model, TauriImageModel::Midjourney | TauriImageModel::Midjourney7 | TauriImageModel::Midjourney7Niji | TauriImageModel::Midjourney8) {
+    return Err(bad_input("This model cannot use a Midjourney account"));
   }
-  
-  Err(GenerateError::ProviderFailure(ProviderFailureReason::MidjourneyJobEnqueueFailed))
+  if has_image_inputs(request) {
+    return Err(bad_input("Direct Midjourney does not support image references yet"));
+  }
+  if request.batch_size.is_some_and(|count| count != 4) {
+    return Err(bad_input("Direct Midjourney generates four images per job"));
+  }
+  let builder = GenerateImageRequestBuilder { model: tauri_image_model_to_router_model(model).unwrap_or(RouterImageModel::Midjourney8), provider: RouterProvider::Midjourney, prompt: request.prompt.clone(), image_inputs: None, resolution: None, aspect_ratio: request.aspect_ratio.map(convert_aspect_ratio), quality: None, image_batch_count: Some(4), horizontal_angle: None, vertical_angle: None, zoom: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, generation_mode_mismatch_strategy: Some(GenerationModeMismatchStrategy::AbortGeneration), idempotency_token: None };
+  let prompt_request = router_image_request_to_artcraft_prompt(&builder);
+  let ImageGenerationDraftOrRequest::Request(generation) = builder.build2()? else {
+    return Err(bad_input("Unexpected Midjourney request draft"));
+  };
+  let session = credentials.session().await.map_err(GenerateError::AnyhowError)?.ok_or_else(GenerateError::needs_midjourney_credentials)?;
+  // Persist the prompt before submit and link it to the task. Completion can
+  // then attribute results after restart, including websocket-only completions.
+  let storyteller_credentials = storyteller.get_credentials().map_err(GenerateError::AnyhowError)?.ok_or_else(GenerateError::needs_storyteller_credentials)?;
+  let prompt = create_prompt(&config.storyteller_host, Some(&storyteller_credentials), prompt_request).await?;
+  let client = RouterClient::Midjourney(RouterMidjourneyClient::new(session.cookie_header, session.user_id, session.browser));
+  let response = generation.send_request(&client).await?;
+  let GenerateImageResponse::Midjourney(response) = response else {
+    return Err(bad_input("Unexpected Midjourney response"));
+  };
+  Ok(TaskEnqueueSuccess { task_type: TaskType::ImageGeneration, model: Some(tauri_image_model_to_generation_model(model)), provider: GenerationProvider::Midjourney, provider_job_id: Some(response.job_id), maybe_queue_status_url: None, maybe_queue_response_url: None, maybe_prompt_token: Some(prompt.prompt_token) })
+}
+
+fn has_image_inputs(request: &TauriGenerateImageRequest) -> bool {
+  request.image_media_tokens.as_ref().is_some_and(|images| !images.is_empty()) || request.canvas_image_media_token.is_some() || request.canvas_image_raw_bytes.is_some() || request.scene_image_media_token.is_some() || request.scene_image_raw_bytes.is_some() || request.inpainting_mask_image_media_token.is_some() || request.inpainting_mask_image_raw_bytes.is_some()
+}
+
+fn bad_input(message: &str) -> GenerateError {
+  GenerateError::BadInput(BadInputReason::WrongImageArguments(message.to_string()))
 }

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/midjourney/mod.rs

@@ -1 +1 @@
-pub (super) mod handle_midjourney;
+pub mod handle_midjourney;

crates/desktop/artcraft/src/core/commands/generate/generate_image/providers/mod.rs

@@ -1,2 +1,3 @@
 pub mod artcraft;
-pub mod artcraft_router;
+pub mod artcraft_router;
+pub mod midjourney;

crates/desktop/artcraft/src/core/commands/generate/generate_image/tauri_generate_image_request.rs

@@ -120,6 +120,7 @@ pub enum TauriGenerateImageErrorType {
   NeedsFalApiKey,
   /// Needs Grok credentials
   NeedsGrokCredentials,
+  NeedsMidjourneyCredentials,
   /// Billing issue
   BillingIssue,
 }

crates/desktop/artcraft/src/core/commands/generate/omni/request.rs

@@ -38,7 +38,7 @@ impl OmniRequest {
   pub fn uses_legacy_image_endpoint(&self) -> bool {
     // These desktop editor integrations still have dedicated request shapes.
     // Catalog models, including any future ID, use Omni by default.
-    matches!(self.model(), Some("recraft_3" | "flux_pro_kontext_max" | "flux_dev_juggernaut" | "flux_pro_1" | "midjourney"))
+    matches!(self.model(), Some("recraft_3" | "flux_pro_kontext_max" | "flux_dev_juggernaut" | "flux_pro_1"))
   }
 
   pub fn api_fields(&self, modality: Modality) -> Map<String, Value> {
@@ -101,6 +101,7 @@ fn rename_field(fields: &mut Map<String, Value>, old: &str, new: &str) {
 // Unknown IDs pass through unchanged; this is not a model allowlist.
 fn canonical_model_id(model: &str) -> &str {
   match model {
+    "midjourney" => "midjourney_8",
     "flux_pro_11" => "flux_pro_1p1",
     "flux_pro_11_ultra" => "flux_pro_1p1_ultra",
     "gemini_25_flash" => "nano_banana",
@@ -185,3 +186,28 @@ mod tests {
     assert_eq!(serde_json::to_value(GenerationModel::FluxPro11).unwrap(), json!("flux_pro_1.1"));
   }
 }
+
+#[cfg(test)]
+mod midjourney_tests {
+  use super::*;
+
+  #[test]
+  fn versioned_models_respect_the_explicit_provider() {
+    for model in ["midjourney_7", "midjourney_7_niji", "midjourney_8"] {
+      let direct: OmniRequest = serde_json::from_value(serde_json::json!({"model": model, "provider": "midjourney"})).unwrap();
+      assert!(!direct.uses_artcraft());
+      let artcraft: OmniRequest = serde_json::from_value(serde_json::json!({"model": model, "provider": "artcraft"})).unwrap();
+      assert!(artcraft.uses_artcraft());
+      assert!(!artcraft.uses_legacy_image_endpoint());
+      assert_eq!(artcraft.api_fields(Modality::Image)["model"], model);
+    }
+  }
+
+  #[test]
+  fn generic_midjourney_uses_v8_when_artcraft_is_selected() {
+    let request: OmniRequest = serde_json::from_value(serde_json::json!({"model": "midjourney", "provider": "artcraft"})).unwrap();
+    assert!(request.uses_artcraft());
+    assert!(!request.uses_legacy_image_endpoint());
+    assert_eq!(request.api_fields(Modality::Image)["model"], "midjourney_8");
+  }
+}

crates/desktop/artcraft/src/core/state/app_preferences/app_preferences.rs

@@ -1,5 +1,6 @@
 use crate::core::state::app_preferences::app_preferences_serializable::AppPreferencesSerializable;
 use crate::core::state::app_preferences::preferred_download_directory::{PreferredDownloadDirectory, SystemDownloadDirectory};
+use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 
 #[derive(Clone)]
@@ -7,6 +8,11 @@ pub struct AppPreferences {
   /// The downloads directory to use when a user downloads a file.
   pub preferred_download_directory: PreferredDownloadDirectory,
 
+  pub preferred_download_filename: PreferredDownloadFilename,
+
+  /// Save generated results before marking their task complete.
+  pub auto_download: bool,
+
   /// Play sounds on events.
   pub play_sounds: bool,
 
@@ -36,6 +42,8 @@ impl Default for AppPreferences {
   fn default() -> Self {
     Self {
       preferred_download_directory: PreferredDownloadDirectory::System(SystemDownloadDirectory::Downloads),
+      preferred_download_filename: PreferredDownloadFilename::default(),
+      auto_download: false,
       play_sounds: true,
       // NB: These are defined in the frontend.
       enqueue_success_sound: Some("done".to_string()),

crates/desktop/artcraft/src/core/state/app_preferences/app_preferences_serializable.rs

@@ -1,5 +1,6 @@
 use crate::core::state::app_preferences::app_preferences::AppPreferences;
 use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory;
+use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use errors::AnyhowResult;
 use serde_derive::{Deserialize, Serialize};
@@ -7,7 +8,8 @@ use serde_derive::{Deserialize, Serialize};
 /// Vector clock versioning string rather than semver.
 /// - Version 1 - initial version.
 /// - Version 2 - added "delete_file_sound", marked optionals "skip_serializing_if"
-const CURRENT_VERSION: &str = "2";
+/// - Version 3 - added auto download and filename convention.
+const CURRENT_VERSION: &str = "3";
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct AppPreferencesSerializable {
@@ -17,6 +19,12 @@ pub struct AppPreferencesSerializable {
   /// The downloads directory to use when a user downloads a file.
   pub preferred_download_directory: Option<PreferredDownloadDirectory>,
 
+  #[serde(rename = "preferred_download_filename", skip_serializing_if = "Option::is_none")]
+  pub maybe_preferred_download_filename: Option<PreferredDownloadFilename>,
+
+  #[serde(rename = "auto_download", skip_serializing_if = "Option::is_none")]
+  pub maybe_auto_download: Option<bool>,
+
   /// Play sounds on events.
   #[serde(skip_serializing_if = "Option::is_none")]
   pub play_sounds: Option<bool>,
@@ -54,6 +62,8 @@ impl AppPreferencesSerializable {
     Self {
       version: CURRENT_VERSION.to_string(),
       preferred_download_directory: Some(preferences.preferred_download_directory.clone()),
+      maybe_preferred_download_filename: Some(preferences.preferred_download_filename.clone()),
+      maybe_auto_download: Some(preferences.auto_download),
       play_sounds: Some(preferences.play_sounds),
       delete_file_sound: preferences.delete_file_sound.clone(),
       generation_success_sound: preferences.generation_success_sound.clone(),
@@ -73,6 +83,13 @@ impl AppPreferencesSerializable {
       preferences.play_sounds = play_sounds;
     }
 
+    if let Some(filename) = &self.maybe_preferred_download_filename {
+      preferences.preferred_download_filename = filename.clone();
+    }
+    if let Some(auto_download) = self.maybe_auto_download {
+      preferences.auto_download = auto_download;
+    }
+
     preferences.delete_file_sound = self.delete_file_sound.clone();
     preferences.generation_success_sound = self.generation_success_sound.clone();
     preferences.generation_failure_sound = self.generation_failure_sound.clone();
@@ -81,3 +98,40 @@ impl AppPreferencesSerializable {
     preferences
   }
 }
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn old_preferences_keep_the_directory_and_default_auto_download_off() {
+    let old: AppPreferencesSerializable = serde_json::from_value(serde_json::json!({
+      "version": "2", "preferred_download_directory": { "custom": "/my/downloads" },
+      "play_sounds": false, "generation_success_sound": "special_flower",
+    })).unwrap();
+    let prefs = old.to_preferences();
+    assert!(!prefs.auto_download);
+    assert_eq!(prefs.preferred_download_filename, PreferredDownloadFilename::ArtcraftConvention);
+    assert_eq!(prefs.preferred_download_directory, PreferredDownloadDirectory::Custom("/my/downloads".into()));
+    assert!(!prefs.play_sounds);
+    assert_eq!(prefs.generation_success_sound.as_deref(), Some("special_flower"));
+  }
+
+  #[test]
+  fn download_preferences_round_trip_without_changing_existing_fields() {
+    let prefs = AppPreferences {
+      auto_download: true,
+      preferred_download_filename: PreferredDownloadFilename::Custom("{model}_{date}".into()),
+      preferred_download_directory: PreferredDownloadDirectory::Custom("/my/downloads".into()),
+      ..AppPreferences::default()
+    };
+    let json = serde_json::to_value(prefs.to_serializable()).unwrap();
+    assert_eq!(json["auto_download"], true);
+    assert_eq!(json["preferred_download_filename"]["custom_format"], "{model}_{date}");
+    let restored = serde_json::from_value::<AppPreferencesSerializable>(json).unwrap().to_preferences();
+    assert!(restored.auto_download);
+    assert_eq!(restored.preferred_download_filename, prefs.preferred_download_filename);
+    assert_eq!(restored.preferred_download_directory, prefs.preferred_download_directory);
+    assert_eq!(restored.generation_success_sound, prefs.generation_success_sound);
+  }
+}

crates/desktop/artcraft/src/core/state/app_preferences/mod.rs

@@ -1,4 +1,5 @@
 pub mod app_preferences;
 pub mod app_preferences_manager;
 pub mod app_preferences_serializable;
-pub mod preferred_download_directory;
+pub mod preferred_download_directory;
+pub mod preferred_download_filename;

crates/desktop/artcraft/src/core/state/app_preferences/preferred_download_filename.rs

@@ -0,0 +1,351 @@
+use chrono::{DateTime, Datelike, Local, Timelike};
+use serde_derive::{Deserialize, Serialize};
+
+/// Characters that may never appear in a download filename (or a custom
+/// filename format): path separators, traversal, quoting, and shell-unsafe
+/// characters.
+const UNSAFE_CHARACTERS: &[char] = &['/', '\\', '\'', '"', '`', '%', '<', '>', '|', ':', '*', '?', '\0'];
+
+/// How downloaded generation files are named on disk.
+///
+/// Stored in the app preferences file. NEVER change existing serialized
+/// values; only add new ones.
+#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
+#[serde(rename_all = "snake_case")]
+pub enum PreferredDownloadFilename {
+  /// `{model}_{date}(_{batch_index}).{ext}`
+  ///  - model slug, eg. "seedance_2.0", "flux_pro_1.1"
+  ///  - date is "YYYY-MM-DD-HH-MM-SS" (download time)
+  ///  - batch index is 1-n when the job produced more than one file
+  ///  - ext matches the cloud file (jpg, png, mp4, ...)
+  ///
+  /// NB: Serializes as `"artcraft_convention"` in JSON.
+  ArtcraftConvention,
+
+  /// A user-defined format string. Supported tokens:
+  ///  - `{model}` — the model slug
+  ///  - `{date}` — "YYYY-MM-DD-HH-MM-SS" (download time)
+  ///  - `{YYYY}`, `{YY}`, `{MM}`, `{DD}`, `{HH}`, `{mm}`, `{SS}` — date parts
+  ///  - `{batch_index}` — 1-n (appended automatically when the job produced
+  ///    more than one file and the format doesn't include it)
+  ///
+  /// Slash, quote, and other unsafe characters are rejected.
+  ///
+  /// NB: Serializes as `{"custom_format": "..."}` in JSON. (The key is NOT
+  /// "custom" so it can't be confused with `PreferredDownloadDirectory`'s
+  /// custom variant in loosely-typed payloads.)
+  #[serde(rename = "custom_format")]
+  Custom(String),
+}
+
+impl Default for PreferredDownloadFilename {
+  fn default() -> Self {
+    Self::ArtcraftConvention
+  }
+}
+
+/// Everything needed to render a download filename.
+pub struct DownloadFilenameParts<'a> {
+  /// Model slug, eg. "seedance_2.0". Use [`model_slug_from_model_type_str`].
+  pub model_slug: &'a str,
+
+  /// The download time.
+  pub download_time: DateTime<Local>,
+
+  /// 1-based index when the job produced more than one file.
+  pub maybe_batch_index: Option<usize>,
+
+  /// File extension without the leading dot, matching the cloud file.
+  pub extension: &'a str,
+}
+
+impl PreferredDownloadFilename {
+  /// Render the filename (including extension) for one downloaded file.
+  /// The output is sanitized: unsafe characters can never appear.
+  pub fn build_filename(&self, parts: &DownloadFilenameParts<'_>) -> String {
+    let stem = match self {
+      Self::ArtcraftConvention => {
+        let mut stem = format!("{}_{}", parts.model_slug, full_date_string(&parts.download_time),);
+        if let Some(index) = parts.maybe_batch_index {
+          stem.push_str(&format!("_{index}"));
+        }
+        stem
+      },
+      Self::Custom(format) => {
+        let mut stem = expand_format_tokens(format, parts);
+        let has_batch_token = format.contains("{batch_index}");
+        if !has_batch_token {
+          if let Some(index) = parts.maybe_batch_index {
+            stem.push_str(&format!("_{index}"));
+          }
+        }
+        stem
+      },
+    };
+
+    // Leave room for an extension and a collision suffix on common filesystems.
+    let mut stem = sanitize(&stem);
+    while stem.len() > 200 {
+      stem.pop();
+    }
+    if stem.is_empty() {
+      stem = "artcraft".to_string();
+    }
+    let device_name = stem.split('.').next().unwrap_or_default().to_ascii_uppercase();
+    if matches!(device_name.as_str(), "CON" | "PRN" | "AUX" | "NUL")
+      || (device_name.len() == 4
+        && (device_name.starts_with("COM") || device_name.starts_with("LPT"))
+        && matches!(device_name.as_bytes()[3], b'1'..=b'9'))
+    {
+      stem.insert_str(0, "artcraft_");
+    }
+    let extension = sanitize(parts.extension);
+    let extension = if extension.is_empty() { "bin" } else { &extension };
+
+    format!("{stem}.{extension}")
+  }
+
+  /// Vet a user-supplied custom format string.
+  pub fn validate_custom_format(format: &str) -> Result<(), String> {
+    if format.trim().is_empty() {
+      return Err("Filename format cannot be empty".to_string());
+    }
+    if format.contains("..") {
+      return Err("Filename format cannot contain '..'".to_string());
+    }
+    if let Some(bad) = format.chars().find(|c| UNSAFE_CHARACTERS.contains(c) || c.is_control()) {
+      return Err(format!("Filename format cannot contain {bad:?}"));
+    }
+    Ok(())
+  }
+}
+
+/// Derive a friendly model slug from a `TaskModelType` string value:
+/// version markers like "2p0" become "2.0" (eg. "seedance_2p0" ->
+/// "seedance_2.0", "flux_pro_1p1" -> "flux_pro_1.1").
+pub fn model_slug_from_model_type_str(model_type: &str) -> String {
+  let mut slug = String::with_capacity(model_type.len());
+  let chars = model_type.chars().collect::<Vec<_>>();
+  for (i, c) in chars.iter().enumerate() {
+    let is_version_dot =
+      *c == 'p' && i > 0 && i + 1 < chars.len() && chars[i - 1].is_ascii_digit() && chars[i + 1].is_ascii_digit();
+    slug.push(if is_version_dot { '.' } else { *c });
+  }
+  slug
+}
+
+fn full_date_string(time: &DateTime<Local>) -> String {
+  format!(
+    "{:04}-{:02}-{:02}-{:02}-{:02}-{:02}",
+    time.year(),
+    time.month(),
+    time.day(),
+    time.hour(),
+    time.minute(),
+    time.second(),
+  )
+}
+
+fn expand_format_tokens(format: &str, parts: &DownloadFilenameParts<'_>) -> String {
+  let time = &parts.download_time;
+  format
+    .replace("{model}", parts.model_slug)
+    .replace("{date}", &full_date_string(time))
+    .replace("{YYYY}", &format!("{:04}", time.year()))
+    .replace("{YY}", &format!("{:02}", time.year() % 100))
+    .replace("{MM}", &format!("{:02}", time.month()))
+    .replace("{DD}", &format!("{:02}", time.day()))
+    .replace("{HH}", &format!("{:02}", time.hour()))
+    .replace("{mm}", &format!("{:02}", time.minute()))
+    .replace("{SS}", &format!("{:02}", time.second()))
+    .replace(
+      "{batch_index}",
+      &parts.maybe_batch_index.map(|i| i.to_string()).unwrap_or_default(),
+    )
+}
+
+/// Strip anything unsafe out of a rendered filename component.
+fn sanitize(value: &str) -> String {
+  value
+    .chars()
+    .filter(|c| !UNSAFE_CHARACTERS.contains(c) && !c.is_control())
+    .collect::<String>()
+    .replace("..", "_")
+    .trim_matches(|c: char| c.is_whitespace() || c == '.')
+    .to_string()
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use chrono::TimeZone;
+
+  mod artcraft_convention {
+    use super::*;
+
+    #[test]
+    fn single_file() {
+      let name = PreferredDownloadFilename::ArtcraftConvention.build_filename(&DownloadFilenameParts {
+        model_slug: "seedance_2.0",
+        download_time: test_time(),
+        maybe_batch_index: None,
+        extension: "mp4",
+      });
+      assert_eq!(name, "seedance_2.0_2026-08-05-14-30-09.mp4");
+    }
+
+    #[test]
+    fn batch_file() {
+      let name = PreferredDownloadFilename::ArtcraftConvention.build_filename(&DownloadFilenameParts {
+        model_slug: "flux_pro_1.1",
+        download_time: test_time(),
+        maybe_batch_index: Some(3),
+        extension: "png",
+      });
+      assert_eq!(name, "flux_pro_1.1_2026-08-05-14-30-09_3.png");
+    }
+  }
+
+  mod custom_format {
+    use super::*;
+
+    #[test]
+    fn tokens_expand() {
+      let format = PreferredDownloadFilename::Custom("{model}-{YYYY}{MM}{DD}-{HH}{mm}{SS}".to_string());
+      let name = format.build_filename(&DownloadFilenameParts {
+        model_slug: "nano_banana",
+        download_time: test_time(),
+        maybe_batch_index: None,
+        extension: "jpg",
+      });
+      assert_eq!(name, "nano_banana-20260805-143009.jpg");
+    }
+
+    #[test]
+    fn batch_index_token() {
+      let format = PreferredDownloadFilename::Custom("gen_{batch_index}".to_string());
+      let name = format.build_filename(&DownloadFilenameParts {
+        model_slug: "unused",
+        download_time: test_time(),
+        maybe_batch_index: Some(2),
+        extension: "png",
+      });
+      assert_eq!(name, "gen_2.png");
+    }
+
+    #[test]
+    fn batch_index_appended_when_absent_from_format() {
+      let format = PreferredDownloadFilename::Custom("gen".to_string());
+      let name = format.build_filename(&DownloadFilenameParts {
+        model_slug: "unused",
+        download_time: test_time(),
+        maybe_batch_index: Some(2),
+        extension: "png",
+      });
+      assert_eq!(name, "gen_2.png");
+    }
+
+    #[test]
+    fn empty_expansions_reserved_names_and_long_unicode_names_are_portable() {
+      let parts = DownloadFilenameParts {
+        model_slug: "midjourney_8",
+        download_time: test_time(),
+        maybe_batch_index: None,
+        extension: "png",
+      };
+      assert_eq!(
+        PreferredDownloadFilename::Custom("{batch_index}".into()).build_filename(&parts),
+        "artcraft.png"
+      );
+      assert_eq!(
+        PreferredDownloadFilename::Custom("CON".into()).build_filename(&parts),
+        "artcraft_CON.png"
+      );
+      let name = PreferredDownloadFilename::Custom("画".repeat(150)).build_filename(&parts);
+      assert!(name.len() < 240);
+      assert!(name.ends_with(".png"));
+    }
+
+    #[test]
+    fn unsafe_characters_stripped_from_output() {
+      let format = PreferredDownloadFilename::Custom("weird{model}".to_string());
+      let name = format.build_filename(&DownloadFilenameParts {
+        model_slug: "a/b\\c\"d",
+        download_time: test_time(),
+        maybe_batch_index: None,
+        extension: "png",
+      });
+      assert_eq!(name, "weirdabcd.png");
+    }
+  }
+
+  mod validation {
+    use super::*;
+
+    #[test]
+    fn accepts_reasonable_formats() {
+      assert!(PreferredDownloadFilename::validate_custom_format("{model}_{date}").is_ok());
+      assert!(PreferredDownloadFilename::validate_custom_format("my file {YYYY}").is_ok());
+    }
+
+    #[test]
+    fn rejects_unsafe_formats() {
+      assert!(PreferredDownloadFilename::validate_custom_format("").is_err());
+      assert!(PreferredDownloadFilename::validate_custom_format("a/b").is_err());
+      assert!(PreferredDownloadFilename::validate_custom_format("a\\b").is_err());
+      assert!(PreferredDownloadFilename::validate_custom_format("a\"b").is_err());
+      assert!(PreferredDownloadFilename::validate_custom_format("a'b").is_err());
+      assert!(PreferredDownloadFilename::validate_custom_format("../etc").is_err());
+    }
+  }
+
+  mod model_slug {
+    use super::*;
+
+    #[test]
+    fn version_markers_become_dots() {
+      assert_eq!(model_slug_from_model_type_str("seedance_2p0"), "seedance_2.0");
+      assert_eq!(model_slug_from_model_type_str("flux_pro_1p1"), "flux_pro_1.1");
+      assert_eq!(model_slug_from_model_type_str("seedream_5p0_pro_u"), "seedream_5.0_pro_u");
+    }
+
+    #[test]
+    fn non_version_names_unchanged() {
+      assert_eq!(model_slug_from_model_type_str("nano_banana"), "nano_banana");
+      assert_eq!(model_slug_from_model_type_str("midjourney_8"), "midjourney_8");
+      assert_eq!(model_slug_from_model_type_str("marble_1p1_plus"), "marble_1.1_plus");
+    }
+  }
+
+  mod serialization {
+    use super::*;
+
+    #[test]
+    fn artcraft_convention_serializes_as_string() {
+      let val = PreferredDownloadFilename::ArtcraftConvention;
+      assert_eq!(serde_json::to_string(&val).unwrap(), "\"artcraft_convention\"");
+    }
+
+    #[test]
+    fn custom_serializes_with_custom_format_key() {
+      let val = PreferredDownloadFilename::Custom("{model}_{date}".to_string());
+      assert_eq!(serde_json::to_string(&val).unwrap(), "{\"custom_format\":\"{model}_{date}\"}");
+    }
+
+    #[test]
+    fn round_trips() {
+      for val in [
+        PreferredDownloadFilename::ArtcraftConvention,
+        PreferredDownloadFilename::Custom("{model}".to_string()),
+      ] {
+        let json = serde_json::to_string(&val).unwrap();
+        let back: PreferredDownloadFilename = serde_json::from_str(&json).unwrap();
+        assert_eq!(back, val);
+      }
+    }
+  }
+
+  fn test_time() -> DateTime<Local> {
+    Local.with_ymd_and_hms(2026, 8, 5, 14, 30, 9).unwrap()
+  }
+}

crates/desktop/artcraft/src/core/state/data_dir/app_credentials_dir.rs

@@ -58,6 +58,10 @@ impl AppCredentialsDir {
     self.path.join("fal_api_key.txt")
   }
 
+  pub fn get_midjourney_credential_path(&self) -> PathBuf {
+    self.path.join("midjourney_cookies.toml")
+  }
+
   pub fn get_midjourney_state_path(&self) -> PathBuf {
     self.path.join("midjourney_state.json")
   }

crates/desktop/artcraft/src/core/threads/third_party_task_polling_thread/handlers/fal/handle_fal_complete.rs

@@ -5,6 +5,7 @@ use crate::core::threads::third_party_task_polling_thread::events::notify_fronte
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
 use crate::core::state::task_database::TaskDatabase;
+use crate::core::utils::auto_download::{auto_download_task_urls, clear_auto_download_checkpoint};
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
 use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
 use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
@@ -89,6 +90,10 @@ async fn handle_fal_complete_inner(
     return Ok(());
   }
 
+  // Save local copies before uploading or announcing completion.
+  let download_urls = urls.iter().map(|url| reqwest::Url::parse(url)).collect::<Result<Vec<_>, _>>()?;
+  auto_download_task_urls(app_handle, task, &download_urls).await?;
+
   // Use the prompt token from the task record (created at generation time).
   let maybe_prompt_token = task.prompt_token.as_ref()
     .map(|s| PromptToken::new_from_str(s));
@@ -163,6 +168,7 @@ async fn handle_fal_complete_inner(
   }).await?;
 
   if updated {
+    clear_auto_download_checkpoint(app_handle, task);
     if let Some(primary_token) = maybe_primary_media_file_token {
       let completion = CompletionData {
         primary_media_file_token: primary_token,
@@ -341,4 +347,3 @@ async fn try_upload(
 
   Ok(media_token)
 }
-

crates/desktop/artcraft/src/core/utils/auto_download.rs

@@ -0,0 +1,309 @@
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::core::state::app_preferences::app_preferences::AppPreferences;
+use crate::core::state::app_preferences::app_preferences_manager::AppPreferencesManager;
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
+use crate::core::utils::download_url_to_user_download_dir::{download_url_with_filename, suggested_download_filename};
+use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
+use anyhow::{anyhow, ensure};
+use artcraft_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
+use chrono::{DateTime, Local};
+use errors::AnyhowResult;
+use serde::{Deserialize, Serialize};
+use sqlite_tasks::queries::task::Task;
+use std::collections::BTreeMap;
+use std::io::Write;
+use std::path::PathBuf;
+use std::time::Duration;
+use tauri::{AppHandle, Manager};
+use tempfile::NamedTempFile;
+use tokens::tokens::batch_generations::BatchGenerationToken;
+use url::Url;
+
+#[derive(Serialize, Deserialize)]
+struct AutoDownloadCheckpoint {
+  download_time: DateTime<Local>,
+  saved: BTreeMap<String, PathBuf>,
+}
+
+/// Await this before updating task status or emitting ANY completion events.
+/// A failed download leaves the task pending for the next polling pass.
+pub async fn auto_download_task(
+  app: &AppHandle,
+  task: &Task,
+  maybe_batch_token: Option<&BatchGenerationToken>,
+  maybe_primary_url: Option<&str>,
+  maybe_expected_count: Option<usize>,
+) -> AnyhowResult<()> {
+  let prefs = app.state::<AppPreferencesManager>().get_clone()?;
+  if !prefs.auto_download {
+    return Ok(());
+  }
+  let urls = if let Some(batch) = maybe_batch_token {
+    let config = app.state::<AppEnvConfigs>();
+    let credentials = app.state::<StorytellerCredentialManager>().get_credentials()?;
+    let result = tokio::time::timeout(
+      Duration::from_secs(60),
+      list_batch_generated_redux_media_files(&config.storyteller_host, credentials.as_ref(), batch),
+    )
+    .await??;
+    ensure!(!result.media_files.is_empty(), "No downloadable results for batch {}", batch);
+    // Stable order across retries, even if the backend reorders the response.
+    let mut files = result.media_files;
+    files.sort_by(|a, b| a.token.as_str().cmp(b.token.as_str()));
+    files.into_iter().map(|file| file.media_links.cdn_url).collect::<Vec<_>>()
+  } else {
+    vec![Url::parse(
+      maybe_primary_url.ok_or_else(|| anyhow!("Completed task has no download URL"))?,
+    )?]
+  };
+  if let Some(expected) = maybe_expected_count {
+    ensure!(
+      urls.len() == expected,
+      "Waiting for all {} download results (currently {})",
+      expected,
+      urls.len()
+    );
+  }
+  download_generated_results(&app.state::<AppDataRoot>(), &prefs, task, &urls).await
+}
+
+/// For providers that already expose all result URLs before upload.
+pub async fn auto_download_task_urls(app: &AppHandle, task: &Task, urls: &[Url]) -> AnyhowResult<()> {
+  let prefs = app.state::<AppPreferencesManager>().get_clone()?;
+  download_generated_results(&app.state::<AppDataRoot>(), &prefs, task, urls).await
+}
+
+pub fn clear_auto_download_checkpoint(app: &AppHandle, task: &Task) {
+  let path = checkpoint_path(&app.state::<AppDataRoot>(), task);
+  if let Err(err) = std::fs::remove_file(path) {
+    if err.kind() != std::io::ErrorKind::NotFound {
+      log::warn!("Could not remove automatic download checkpoint: {}", err);
+    }
+  }
+}
+
+async fn download_generated_results(
+  root: &AppDataRoot,
+  prefs: &AppPreferences,
+  task: &Task,
+  urls: &[Url],
+) -> AnyhowResult<()> {
+  if !prefs.auto_download {
+    return Ok(());
+  }
+  ensure!(!urls.is_empty(), "Completed task has no downloadable results");
+  let path = checkpoint_path(root, task);
+  let mut checkpoint: AutoDownloadCheckpoint = match std::fs::read(&path) {
+    Ok(bytes) => serde_json::from_slice(&bytes)?,
+    Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
+      AutoDownloadCheckpoint { download_time: Local::now(), saved: BTreeMap::new() }
+    },
+    Err(err) => return Err(err.into()),
+  };
+  let directory = prefs.preferred_download_directory.download_directory(root);
+  for (index, url) in urls.iter().enumerate() {
+    if checkpoint.saved.get(url.as_str()).is_some_and(|path| path.is_file()) {
+      continue;
+    }
+    let filename = suggested_download_filename(
+      url,
+      prefs,
+      task.model_type.as_ref().map(|model| model.to_str()),
+      (urls.len() > 1).then_some(index + 1),
+      checkpoint.download_time,
+    );
+    let saved = download_url_with_filename(url, &directory, &filename).await?;
+    checkpoint.saved.insert(url.to_string(), saved);
+    let mut file = NamedTempFile::new_in(root.state_dir().path())?;
+    file.write_all(&serde_json::to_vec(&checkpoint)?)?;
+    file.as_file().sync_all()?;
+    file.persist(&path)?;
+  }
+  Ok(())
+}
+
+fn checkpoint_path(root: &AppDataRoot, task: &Task) -> PathBuf {
+  root.state_dir().path().join(format!("auto_download_{}.json", task.id.as_str()))
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use crate::core::state::app_preferences::preferred_download_directory::PreferredDownloadDirectory;
+  use crate::core::state::app_preferences::preferred_download_filename::PreferredDownloadFilename;
+  use enums::common::generation_provider::GenerationProvider;
+  use enums::tauri::tasks::{task_model_type::TaskModelType, task_status::TaskStatus, task_type::TaskType};
+  use tokio::io::{AsyncReadExt, AsyncWriteExt};
+  use tokio::net::TcpListener;
+  use tokio::sync::oneshot;
+  use tokens::tokens::sqlite::tasks::TaskId;
+
+  #[tokio::test]
+  async fn disabled_auto_download_does_not_access_network_or_create_files() {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    let mut prefs = preferences(&root);
+    prefs.auto_download = false;
+    download_generated_results(
+      &root,
+      &prefs,
+      &task(GenerationProvider::Artcraft),
+      &[Url::parse("http://127.0.0.1:1/never.png").unwrap()],
+    )
+    .await
+    .unwrap();
+    assert!(!dir.path().join("chosen").exists());
+  }
+
+  #[tokio::test]
+  async fn all_results_are_saved_for_artcraft_and_midjourney_and_retries_skip_saved_files() {
+    for provider in [GenerationProvider::Artcraft, GenerationProvider::Midjourney] {
+      let dir = tempfile::tempdir().unwrap();
+      let root = AppDataRoot::create_existing(dir.path()).unwrap();
+      let prefs = preferences(&root);
+      let task = task(provider);
+      let (base, server) = server(vec![(200, "image"), (200, "video"), (200, "audio"), (200, "splat")]).await;
+      let urls = ["png", "mp4", "mp3", "spz"].map(|ext| Url::parse(&format!("{base}/result.{ext}")).unwrap());
+      download_generated_results(&root, &prefs, &task, &urls).await.unwrap();
+      server.await.unwrap();
+      for (index, (ext, bytes)) in [("png", "image"), ("mp4", "video"), ("mp3", "audio"), ("spz", "splat")]
+        .iter()
+        .enumerate()
+      {
+        assert_eq!(
+          std::fs::read_to_string(dir.path().join(format!("chosen/midjourney_8_result_{}.{ext}", index + 1))).unwrap(),
+          *bytes
+        );
+      }
+      // The server is gone; resuming from disk must not redownload anything.
+      download_generated_results(&root, &prefs, &task, &urls).await.unwrap();
+      assert_eq!(std::fs::read_dir(dir.path().join("chosen")).unwrap().count(), 4);
+    }
+  }
+
+  #[tokio::test]
+  async fn partial_failure_keeps_finished_files_and_retries_only_the_failed_result() {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    let prefs = preferences(&root);
+    let task = task(GenerationProvider::Midjourney);
+    let (base, server) = server(vec![(200, "first"), (500, "error"), (200, "second")]).await;
+    let urls = ["first", "second"].map(|name| Url::parse(&format!("{base}/{name}.png")).unwrap());
+    assert!(download_generated_results(&root, &prefs, &task, &urls).await.is_err());
+    assert_eq!(std::fs::read_dir(dir.path().join("chosen")).unwrap().count(), 1);
+    download_generated_results(&root, &prefs, &task, &urls).await.unwrap();
+    server.await.unwrap();
+    assert_eq!(
+      std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result_1.png")).unwrap(),
+      "first"
+    );
+    assert_eq!(
+      std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result_2.png")).unwrap(),
+      "second"
+    );
+  }
+
+  #[tokio::test]
+  async fn matching_custom_filenames_never_overwrite_previous_generations() {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    let prefs = preferences(&root);
+    let (base, server) = server(vec![(200, "first"), (200, "second")]).await;
+    let url = Url::parse(&format!("{base}/file.png")).unwrap();
+    for _ in 0..2 {
+      download_generated_results(&root, &prefs, &task(GenerationProvider::Artcraft), &[url.clone()])
+        .await
+        .unwrap();
+    }
+    server.await.unwrap();
+    assert_eq!(
+      std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result.png")).unwrap(),
+      "first"
+    );
+    assert_eq!(
+      std::fs::read_to_string(dir.path().join("chosen/midjourney_8_result_1.png")).unwrap(),
+      "second"
+    );
+  }
+
+  #[tokio::test]
+  async fn completion_waits_until_the_entire_file_is_on_disk() {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    let prefs = preferences(&root);
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let url = Url::parse(&format!("http://{}/file.png", listener.local_addr().unwrap())).unwrap();
+    let (started_tx, started_rx) = oneshot::channel();
+    let (finish_tx, finish_rx) = oneshot::channel();
+    let server = tokio::spawn(async move {
+      let (mut socket, _) = listener.accept().await.unwrap();
+      let mut request = [0; 4096];
+      socket.read(&mut request).await.unwrap();
+      socket
+        .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 6\r\nConnection: close\r\n\r\nabc")
+        .await
+        .unwrap();
+      started_tx.send(()).unwrap();
+      finish_rx.await.unwrap();
+      socket.write_all(b"def").await.unwrap();
+    });
+    let completion = tokio::spawn(async move {
+      download_generated_results(&root, &prefs, &task(GenerationProvider::Artcraft), &[url])
+        .await
+        .unwrap();
+      // Callers may mark success and emit notifications only after this await.
+      std::fs::read_to_string(root.path().join("chosen/midjourney_8_result.png")).unwrap()
+    });
+    started_rx.await.unwrap();
+    assert!(!completion.is_finished());
+    assert!(!dir.path().join("chosen/midjourney_8_result.png").exists());
+    finish_tx.send(()).unwrap();
+    assert_eq!(completion.await.unwrap(), "abcdef");
+    server.await.unwrap();
+  }
+
+  fn preferences(root: &AppDataRoot) -> AppPreferences {
+    AppPreferences {
+      auto_download: true,
+      preferred_download_directory: PreferredDownloadDirectory::Custom(root.path().join("chosen")),
+      preferred_download_filename: PreferredDownloadFilename::Custom("{model}_result".into()),
+      ..AppPreferences::default()
+    }
+  }
+
+  fn task(provider: GenerationProvider) -> Task {
+    Task {
+      id: TaskId::generate(),
+      status: TaskStatus::Pending,
+      task_type: TaskType::ImageGeneration,
+      model_type: Some(TaskModelType::Midjourney8),
+      provider,
+      provider_job_id: None,
+      queue_status_url: None,
+      queue_response_url: None,
+      prompt_token: None,
+      frontend_caller: None,
+      frontend_subscriber_id: None,
+      frontend_subscriber_payload: None,
+    }
+  }
+
+  async fn server(responses: Vec<(u16, &'static str)>) -> (String, tokio::task::JoinHandle<()>) {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let base = format!("http://{}", listener.local_addr().unwrap());
+    let server = tokio::spawn(async move {
+      for (status, body) in responses {
+        let (mut socket, _) = listener.accept().await.unwrap();
+        let mut request = [0; 4096];
+        socket.read(&mut request).await.unwrap();
+        let response = format!(
+          "HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
+          body.len()
+        );
+        socket.write_all(response.as_bytes()).await.unwrap();
+      }
+    });
+    (base, server)
+  }
+}

crates/desktop/artcraft/src/core/utils/download_url_to_user_download_dir.rs

@@ -1,12 +1,15 @@
 use crate::core::artcraft_error::ArtcraftError;
 use crate::core::state::app_preferences::app_preferences::AppPreferences;
+use crate::core::state::app_preferences::preferred_download_filename::{model_slug_from_model_type_str, DownloadFilenameParts, PreferredDownloadFilename};
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use anyhow::anyhow;
 use errors::AnyhowResult;
 use log::{error, info};
-use std::fs::OpenOptions;
-use std::io::Write;
-use std::path::PathBuf;
+use chrono::{DateTime, Local};
+use std::path::{Path, PathBuf};
+use std::time::Duration;
+use tempfile::NamedTempFile;
+use tokio::io::AsyncWriteExt;
 use url::Url;
 
 pub async fn download_url_to_user_download_dir(
@@ -14,67 +17,64 @@ pub async fn download_url_to_user_download_dir(
   app_data_root: &AppDataRoot,
   app_prefs: &AppPreferences,
 ) -> Result<PathBuf, ArtcraftError> {
-
-  let url_file_name = {
-    let path_segments = url.path_segments()
-        .ok_or_else(|| ArtcraftError::AnyhowError(anyhow!("URL does not have path segments (1): {:?}", url)))?;
-
-    let mut url_file_name = path_segments.last()
-        .ok_or_else(|| ArtcraftError::AnyhowError(anyhow!("URL does not have path segments (2): {:?}", url)))?
-        .trim()
-        .to_string();
-
-    if !url_file_name.to_lowercase().starts_with("artcraft") {
-      url_file_name = format!("artcraft_{}", url_file_name);
-    }
-
-    url_file_name
-  };
-
-  check_url_file_name_for_downloadability(&url_file_name)?;
-
+  let filename = suggested_download_filename(url, app_prefs, None, None, Local::now());
   let download_directory = app_prefs
       .preferred_download_directory
       .download_directory(app_data_root);
+  download_url_with_filename(url, &download_directory, &filename).await
+}
 
-  info!("Pushing path `{}` onto downloaded directory `{:?}`", url_file_name, download_directory);
+pub fn suggested_download_filename(
+  url: &Url,
+  prefs: &AppPreferences,
+  maybe_model: Option<&str>,
+  maybe_batch_index: Option<usize>,
+  download_time: DateTime<Local>,
+) -> String {
+  let extension = Path::new(url.path()).extension().and_then(|ext| ext.to_str()).unwrap_or("bin");
+  let model_slug = model_slug_from_model_type_str(maybe_model.unwrap_or("artcraft"));
+  prefs.preferred_download_filename.build_filename(&DownloadFilenameParts {
+    model_slug: &model_slug,
+    download_time,
+    maybe_batch_index,
+    extension,
+  })
+}
 
-  let download_filename = {
-    let mut download_filename = download_directory.clone();
-    download_filename.push(&url_file_name);
-    download_filename
-  };
-
-  info!("Download filename: `{:?}`", download_filename);
-  
-  if download_filename == download_directory {
-    return Err(ArtcraftError::AnyhowError(anyhow!("Download filename resolved to directory: {:?}", download_filename)))?;
+/// Stream to a temporary file in the destination directory, then publish the
+/// complete file atomically. Never overwrite another download, even if two
+/// generations finish in the same second or use the same custom name.
+pub async fn download_url_with_filename(url: &Url, directory: &Path, filename: &str) -> Result<PathBuf, ArtcraftError> {
+  check_url_file_name_for_downloadability(filename)?;
+  tokio::fs::create_dir_all(directory).await?;
+  let temporary = NamedTempFile::new_in(directory)?;
+  let mut file = tokio::fs::File::from_std(temporary.reopen()?);
+  let client = reqwest::Client::builder().connect_timeout(Duration::from_secs(30)).timeout(Duration::from_secs(600)).build()?;
+  let mut response = client.get(url.clone()).send().await?.error_for_status()?;
+  while let Some(chunk) = response.chunk().await? {
+    file.write_all(&chunk).await?;
   }
+  file.flush().await?;
+  file.sync_all().await?;
+  drop(file);
+  let path = persist_download(temporary, directory, filename)?;
+  info!("Downloaded generation to {:?}", path);
+  Ok(path)
+}
 
-  if download_filename.exists() {
-    if download_filename.is_dir() {
-      return Err(ArtcraftError::AnyhowError(anyhow!("Download path exists and resolved to directory: {:?}", download_filename)));
-    } else {
-      return Err(ArtcraftError::CannotDownloadFilePathAlreadyExists { path: download_filename }.into());
+fn persist_download(mut file: NamedTempFile, directory: &Path, filename: &str) -> Result<PathBuf, ArtcraftError> {
+  let stem = Path::new(filename).file_stem().and_then(|stem| stem.to_str()).unwrap_or("artcraft");
+  let extension = Path::new(filename).extension().and_then(|ext| ext.to_str()).unwrap_or("bin");
+  for suffix in 0..10_000 {
+    let name = if suffix == 0 { filename.to_string() } else { format!("{stem}_{suffix}.{extension}") };
+    let path = directory.join(name);
+    match file.persist_noclobber(&path) {
+      Ok(_) => return Ok(path),
+      Err(err) if err.error.kind() == std::io::ErrorKind::AlreadyExists => file = err.file,
+      Err(err) => return Err(err.error.into()),
     }
   }
-
-  info!("Downloading bytes from url: {:?}", url);
-
-  let response = reqwest::get(url.clone()).await?;
-  let response_bytes = response.bytes().await?;
-
-  info!("Writing file: {:?}", download_filename);
-
-  let mut file = OpenOptions::new()
-      .create(true)
-      .write(true)
-      .truncate(true)
-      .open(&download_filename)?;
-
-  file.write_all(&response_bytes)?;
-
-  Ok(download_filename)
+  Err(anyhow!("Too many downloads with the same filename").into())
 }
 
 // TODO: This likely needs more protections.
@@ -84,10 +84,7 @@ fn check_url_file_name_for_downloadability(filename: &str) -> Result<(), Artcraf
     return Err(ArtcraftError::BadDownloadFilename { path: "".into() }.into());
   }
 
-  if !filename.is_ascii() {
-    error!("Download filename is not completely ASCII: {:?}", filename);
-    return Err(ArtcraftError::BadDownloadFilename { path: filename.into() }.into());
-  }
+  PreferredDownloadFilename::validate_custom_format(filename).map_err(|err| anyhow!(err))?;
 
   if filename.contains("/")
       || filename.contains("..")
@@ -108,6 +105,8 @@ fn check_url_file_name_for_downloadability(filename: &str) -> Result<(), Artcraf
     return Err(ArtcraftError::BadDownloadFilename { path: filename.into() }.into());
   }
 
+  let filename_lower = filename.to_ascii_lowercase();
+  let filename = filename_lower.as_str();
   if filename.ends_with(".apk")
       || filename.ends_with(".app")
       || filename.ends_with(".bat")

crates/desktop/artcraft/src/core/utils/mod.rs

@@ -1,4 +1,5 @@
 pub mod clear_all_webview_cookies;
+pub mod auto_download;
 pub mod download_media_file_to_temp_dir;
 pub mod download_url_to_temp_dir;
 pub mod download_url_to_user_download_dir;
@@ -12,4 +13,4 @@ pub mod simple_http_download;
 pub mod simple_http_download_to_tempfile;
 pub mod task_database_pending_statuses;
 pub mod upload_bytes_to_media_file;
-pub mod window;
+pub mod window;

crates/desktop/artcraft/src/lib.rs

@@ -19,6 +19,7 @@ use crate::core::commands::generate::models::video::list_video_models_command::l
 use crate::core::commands::download::download_directory_reveal_command::download_directory_reveal_command;
 use crate::core::commands::download::download_media_file_command::download_media_file_command;
 use crate::core::commands::download::download_url_command::download_url_command;
+use crate::core::commands::download::get_download_path_command::get_download_path_command;
 use crate::core::commands::enqueue::image_bg_removal::enqueue_image_bg_removal_command::enqueue_image_bg_removal_command;
 use crate::core::commands::enqueue::image_to_gaussian::enqueue_image_to_gaussian_command::enqueue_image_to_gaussian_command;
 use crate::core::commands::enqueue::image_to_object::enqueue_image_to_3d_object_command::enqueue_image_to_3d_object_command;
@@ -220,6 +221,7 @@ pub fn run() {
     download_directory_reveal_command,
     download_media_file_command,
     download_url_command,
+    get_download_path_command,
     enqueue_image_bg_removal_command,
     enqueue_image_to_3d_object_command,
     enqueue_image_to_gaussian_command,

crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs

@@ -2,6 +2,7 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
 use crate::core::events::functional_events::text_to_image_generation_complete_event::{GeneratedImage, TextToImageGenerationCompleteEvent};
 use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
 use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint};
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::task_database::TaskDatabase;
@@ -248,8 +249,6 @@ async fn upload_images_to_storyteller(
       model: None,
     };
 
-    event.send_infallible(&app_handle);
-
     let task = get_task_by_provider_and_provider_job_id(GetTaskByProviderAndProviderJobIdArgs {
       db: task_database.get_connection(),
       provider: GenerationProvider::Grok,
@@ -286,6 +285,8 @@ async fn upload_images_to_storyteller(
       }
     }
 
+    auto_download_task(app_handle, &task, Some(&batch_token), maybe_cdn_url.as_deref(), Some(images.images.len())).await?;
+
     let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
       db: task_database.get_connection(),
       task_id: &task.id,
@@ -300,6 +301,8 @@ async fn upload_images_to_storyteller(
       return Ok(()); // If anything breaks with queries, don't spam events.
     }
 
+    clear_auto_download_checkpoint(app_handle, &task);
+    event.send_infallible(app_handle);
     send_frontend_ui_update(
       app_handle,
       app_env_configs,

crates/desktop/artcraft/src/services/grok/threads/grok_video_task_polling/grok_video_task_polling_thread.rs

@@ -1,6 +1,7 @@
 use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
 use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
 use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint};
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
@@ -326,6 +327,8 @@ async fn upload_grok_video(
     }
   }
 
+  auto_download_task(app_handle, local_task, None, maybe_cdn_url.as_deref(), None).await?;
+
   let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
     db: task_database.get_connection(),
     task_id: &local_task.id,
@@ -340,6 +343,7 @@ async fn upload_grok_video(
     return Ok(()); // If anything breaks with queries, don't spam events.
   }
 
+  clear_auto_download_checkpoint(app_handle, local_task);
   let event = GenerationCompleteEvent {
     //media_file_token: result.media_file_token,
     action: Some(GenerationAction::GenerateVideo),
@@ -365,4 +369,3 @@ async fn upload_grok_video(
 
   Ok(())
 }
-

crates/desktop/artcraft/src/services/midjourney/commands/midjourney_clear_credentials_command.rs

@@ -1,39 +1,16 @@
-use crate::core::commands::response::shorthand::{Response, SimpleResponse};
-use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::commands::response::shorthand::SimpleResponse;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
-use errors::AnyhowResult;
 use log::{error, info};
-use serde_derive::Deserialize;
 use tauri::State;
 
 #[tauri::command]
-pub async fn midjourney_clear_credentials_command(
-  root: State<'_, AppDataRoot>,
-  creds_manager: State<'_, MidjourneyCredentialManager>,
-) -> SimpleResponse {
+pub async fn midjourney_clear_credentials_command(creds_manager: State<'_, MidjourneyCredentialManager>) -> SimpleResponse {
   info!("midjourney_clear_credentials_command called");
 
-  clear_creds(&root, &creds_manager)
-      .map_err(|err| {
-        error!("Error clearing creds: {:?}", err);
-        "error clearing creds"
-      })?;
+  creds_manager.clear_credentials().await.map_err(|err| {
+    error!("Error clearing creds: {:?}", err);
+    "error clearing creds"
+  })?;
 
   Ok(().into())
 }
-
-fn clear_creds(
-  root: &AppDataRoot,
-  creds: &MidjourneyCredentialManager,
-) -> AnyhowResult<()> {
-
-  creds.clear_credentials()?;
-  creds.persist_to_disk()?;
-
-  let creds_path = root.credentials_dir().get_midjourney_state_path();
-  if creds_path.exists() {
-    std::fs::remove_file(creds_path)?;
-  }
-  
-  Ok(())
-}

crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential.rs

@@ -0,0 +1,84 @@
+use browser_emulation::browser_profile::BrowserProfile;
+use chrono::{DateTime, Utc};
+use cookie_store_wrapper::cookie_store::CookieStore;
+use serde::{Deserialize, Serialize};
+use std::io::Write;
+use std::path::Path;
+use tempfile::NamedTempFile;
+use uuid_utils::uuid::generate_random_uuid;
+
+use crate::services::midjourney::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
+
+/// Midjourney's slice of ArtCraftX's credential format. Other accounts retain
+/// their existing storage. Websocket tokens are deliberately never persisted.
+#[derive(Clone, Serialize, Deserialize)]
+pub struct MidjourneyCredential {
+  #[serde(default = "new_credential_id")]
+  pub id: String,
+  pub service: MidjourneyService,
+  #[serde(default, rename = "name", skip_serializing_if = "Option::is_none")]
+  pub maybe_name: Option<String>,
+  #[serde(default, rename = "user_info", skip_serializing_if = "Option::is_none")]
+  pub maybe_user_info: Option<MidjourneyAccountInfo>,
+  pub cookie: MidjourneyCookieCredential,
+}
+
+#[derive(Clone, Serialize, Deserialize)]
+pub enum MidjourneyService {
+  #[serde(rename = "midjourney_cookies")]
+  Cookies,
+}
+
+#[derive(Clone, Default, Serialize, Deserialize)]
+pub struct MidjourneyAccountInfo {
+  #[serde(default, rename = "user_id", skip_serializing_if = "Option::is_none")]
+  pub maybe_user_id: Option<String>,
+  #[serde(default, rename = "email", skip_serializing_if = "Option::is_none")]
+  pub maybe_email: Option<String>,
+  #[serde(default, rename = "username", skip_serializing_if = "Option::is_none")]
+  pub maybe_username: Option<String>,
+}
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct MidjourneyCookieCredential {
+  #[serde(default, rename = "updated_at", skip_serializing_if = "Option::is_none")]
+  pub maybe_updated_at: Option<DateTime<Utc>>,
+  #[serde(default, rename = "failed_at", skip_serializing_if = "Option::is_none")]
+  pub maybe_failed_at: Option<DateTime<Utc>>,
+  #[serde(default, rename = "succeeded_at", skip_serializing_if = "Option::is_none")]
+  pub maybe_succeeded_at: Option<DateTime<Utc>>,
+  #[serde(default, rename = "relogin_required_since", skip_serializing_if = "Option::is_none")]
+  pub maybe_relogin_required_since: Option<DateTime<Utc>>,
+  #[serde(default, rename = "user_agent", skip_serializing_if = "Option::is_none")]
+  pub maybe_user_agent: Option<String>,
+  // TOML arrays of tables belong after scalar fields.
+  pub cookies: CookieStore,
+}
+
+impl MidjourneyCredential {
+  pub fn new(cookies: CookieStore) -> Self {
+    Self { id: new_credential_id(), service: MidjourneyService::Cookies, maybe_name: None, maybe_user_info: None, cookie: MidjourneyCookieCredential { maybe_updated_at: Some(Utc::now()), maybe_failed_at: None, maybe_succeeded_at: None, maybe_relogin_required_since: None, maybe_user_agent: Some(MIDJOURNEY_USER_AGENT.to_string()), cookies } }
+  }
+
+  pub fn load(path: &Path) -> anyhow::Result<Self> {
+    Ok(toml::from_str(&std::fs::read_to_string(path)?)?)
+  }
+
+  pub fn save(&self, path: &Path) -> anyhow::Result<()> {
+    // Atomic replacement avoids truncating the only session on a failed write.
+    // NamedTempFile creates the file with private permissions on Unix.
+    let mut file = NamedTempFile::new_in(path.parent().ok_or_else(|| anyhow::anyhow!("Missing credential directory"))?)?;
+    file.write_all(toml::to_string_pretty(self)?.as_bytes())?;
+    file.as_file().sync_all()?;
+    file.persist(path)?;
+    Ok(())
+  }
+
+  pub fn browser(&self) -> BrowserProfile {
+    BrowserProfile::matching_user_agent(self.cookie.maybe_user_agent.as_deref().unwrap_or(MIDJOURNEY_USER_AGENT))
+  }
+}
+
+fn new_credential_id() -> String {
+  format!("credential_{}", generate_random_uuid())
+}

crates/desktop/artcraft/src/services/midjourney/state/midjourney_credential_manager.rs

@@ -1,196 +1,352 @@
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::services::midjourney::state::midjourney_credential::{MidjourneyAccountInfo, MidjourneyCredential};
 use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
-use crate::services::midjourney::state::serializable_midjourney_state::{SerializableMidjourneyState, SERIALIZABLE_MIDJOURNEY_STATE_VERSION};
-use crate::services::storyteller::state::read_storyteller_credentials_from_disk::read_storyteller_credentials_from_disk;
-use crate::services::storyteller::state::storyteller_credential_holder::StorytellerCredentialHolder;
-use cookie_store::cookie_store::CookieStore;
-use errors::AnyhowResult;
-use log::{info, warn};
-use std::fs::read_to_string;
-use std::sync::{Arc, RwLock};
+use crate::services::midjourney::state::serializable_midjourney_state::SerializableMidjourneyState;
+use crate::services::midjourney::utils::extract_midjourney_user_id_from_cookies::extract_midjourney_user_id_from_cookie_header;
+use anyhow::{anyhow, bail};
+use base64::prelude::{Engine, BASE64_URL_SAFE_NO_PAD};
+use browser_emulation::browser_profile::BrowserProfile;
+use chrono::Utc;
+use cookie_store_wrapper::cookie_store::CookieStore;
 use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
-use artcraft_client::credentials::storyteller_avt_cookie::StorytellerAvtCookie;
-use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
-use artcraft_client::credentials::storyteller_session_cookie::StorytellerSessionCookie;
+use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
+use midjourney_client::endpoints::refresh_token::{refresh_token, RefreshTokenArgs, RefreshTokenRequest, DEFAULT_FIREBASE_API_KEY};
+use midjourney_client::error::midjourney_api_error::MidjourneyApiError;
+use midjourney_client::error::midjourney_error::MidjourneyError;
+use midjourney_client::recipes::get_user_info::{get_user_info, GetUserInfoArgs};
+use serde_json::Value;
+use std::sync::{Arc, RwLock};
+use std::time::Duration;
+use tokio::sync::Mutex;
+use url::Url;
+
+const AUTH_ID: &str = "__Host-Midjourney.AuthUserTokenV3_i";
+const AUTH_REFRESH: &str = "__Host-Midjourney.AuthUserTokenV3_r";
+const MIDJOURNEY_ORIGIN: &str = "https://www.midjourney.com/";
 
 #[derive(Clone)]
 pub struct MidjourneyCredentialManager {
-  // TODO: Put last write/last read timestamps on these.
-  cookies: Arc<RwLock<Option<CookieStore>>>,
-  user_info: Arc<RwLock<Option<MidjourneyUserInfo>>>,
+  credential: Arc<RwLock<Option<MidjourneyCredential>>>,
+  // Serialize refresh, login, and logout so an old refresh cannot resurrect a
+  // signed-out session or overwrite credentials captured by a new login.
+  mutation: Arc<Mutex<()>>,
   app_data_root: AppDataRoot,
 }
 
+pub struct MidjourneySession {
+  pub credential_id: String,
+  pub cookie_header: String,
+  pub user_id: MidjourneyUserId,
+  pub browser: BrowserProfile,
+}
+
 impl MidjourneyCredentialManager {
   pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self {
-    Self {
-      cookies: Arc::new(RwLock::new(None)),
-      user_info: Arc::new(RwLock::new(None)),
-      app_data_root: app_data_root.clone(),
-    }
+    Self { credential: Arc::new(RwLock::new(None)), mutation: Arc::new(Mutex::new(())), app_data_root: app_data_root.clone() }
   }
 
   pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self {
-    let mut cookies;
-    let mut user_info;
-    
-    match read_midjourney_state_from_disk(app_data_root) {
-      Err(err) => {
-        warn!("Failed to read midjourney state from disk: {:?}", err);
-        cookies = Arc::new(RwLock::new(None));
-        user_info = Arc::new(RwLock::new(None));
+    let manager = Self::initialize_empty(app_data_root);
+    match load_or_migrate(app_data_root) {
+      Ok(credential) => {
+        if let Ok(mut current) = manager.credential.write() {
+          *current = credential;
+        }
       },
-      Ok(None) => {
-        cookies = Arc::new(RwLock::new(None));
-        user_info = Arc::new(RwLock::new(None));
-      }
-      Ok(Some(state)) => {
-        let maybe_cookies = state.user_cookies
-            .map(|cookies| cookies.to_cookie_store());
-        let maybe_user_info = state.user_info
-            .map(|info| info.to_user_info());
-        cookies = Arc::new(RwLock::new(maybe_cookies));
-        user_info = Arc::new(RwLock::new(maybe_user_info));
-      }
-    };
-
-    Self {
-      cookies,
-      user_info,
-      app_data_root: app_data_root.clone(),
+      Err(err) => log::warn!("Could not load Midjourney credentials: {}", err),
     }
+    manager
   }
-  
+
   pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
-    match self.cookies.read() {
-      Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
-      Ok(store) => Ok(store.clone()),
-    }
+    Ok(self.snapshot()?.map(|credential| credential.cookie.cookies))
   }
-  
+
   pub fn maybe_copy_user_info(&self) -> anyhow::Result<Option<MidjourneyUserInfo>> {
-    match self.user_info.read() {
-      Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
-      Ok(info) => Ok(info.clone()),
-    }
-  }
-  
-  pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> {
-    match self.cookies.write() {
-      Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
-      Ok(mut current_store) => {
-        *current_store = Some(store);
-        Ok(())
-      }
-    }
+    Ok(self.snapshot()?.and_then(|credential| credential.maybe_user_info).map(|info| MidjourneyUserInfo { user_id: info.maybe_user_id.map(MidjourneyUserId::from_string), email: info.maybe_email, websocket_token: None }))
   }
 
-  pub fn replace_user_info(&self, user_info: MidjourneyUserInfo) -> anyhow::Result<()> {
-    match self.user_info.write() {
-      Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
-      Ok(mut current_info) => {
-        *current_info = Some(user_info);
-        Ok(())
-      }
-    }
-  }
-
-  // NB: This is just a heuristic. We'll add better checks later.
   pub fn session_appears_active(&self) -> anyhow::Result<bool> {
-    let maybe_cookies = match self.cookies.read() {
-      Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
-      Ok(store) => store.clone(),
-    };
-
-    let maybe_user_info = match self.user_info.read() {
-      Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
-      Ok(info) => info.clone(),
-    };
-
-    if maybe_cookies.is_none() || maybe_user_info.is_none() {
-      return Ok(false);
-    }
-
-    let cookies = maybe_cookies.unwrap();
-    let user_info = maybe_user_info.unwrap();
-
-    // NB: We consider the session active if we have auth cookies and a user id.
-    let has_user_id = user_info.user_id.is_some();
-    let maybe_has_auth_cookies = cookie_store_has_auth_cookies(&cookies);
-
-    // Misc cookies without login cookies are ~1055 length
-    // AUTH_I is ~1500 length
-    // AUTH_R is ~500 length
-    let maybe_has_big_enough_cookie = cookies.calculate_approx_cookie_character_length() > 2100;
-
-    // TODO: Heuristic should count.
-    // TODO: Consolidate with "login window thread" logic.
-    // TODO: Check timestamp of last valid request.
-    Ok(has_user_id && maybe_has_auth_cookies && maybe_has_big_enough_cookie && maybe_has_big_enough_cookie)
+    Ok(self.snapshot()?.is_some_and(|credential| credential.cookie.maybe_relogin_required_since.is_none() && (cookie_store_has_auth_cookies(&credential.cookie.cookies) || !token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_ID), Utc::now().timestamp()))))
   }
 
-  pub fn clear_credentials(&self) -> anyhow::Result<()> {
-    match self.cookies.write() {
-      Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
-      Ok(mut store) => *store = None,
-    }
-    match self.user_info.write() {
-      Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
-      Ok(mut info) => *info = None,
+  pub async fn save_login(&self, cookies: CookieStore) -> anyhow::Result<()> {
+    let _guard = self.mutation.lock().await;
+    let mut credential = MidjourneyCredential::new(cookies);
+    populate_account_info(&mut credential);
+    self.store(credential)
+  }
+
+  pub async fn clear_credentials(&self) -> anyhow::Result<()> {
+    let _guard = self.mutation.lock().await;
+    // Remove both formats, so a later startup cannot migrate a logged-out account.
+    for path in [self.app_data_root.credentials_dir().get_midjourney_state_path(), self.credential_path()] {
+      match std::fs::remove_file(path) {
+        Ok(()) => {},
+        Err(err) if err.kind() == std::io::ErrorKind::NotFound => {},
+        Err(err) => return Err(err.into()),
+      }
     }
+    *self.credential.write().map_err(|_| anyhow!("Midjourney credential lock poisoned"))? = None;
     Ok(())
   }
-  
-  pub fn persist_to_disk(&self) -> anyhow::Result<()> {
-    let maybe_cookies = match self.cookies.read() {
-      Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
-      Ok(store) => store.clone(),
-    };
 
-    let maybe_user_info = match self.user_info.read() {
-      Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
-      Ok(info) => info.clone(),
+  /// Resolve identity and renew expiring Firebase credentials before sending
+  /// generation or polling requests. The refresh token can outlive the ID JWT.
+  pub async fn session(&self) -> anyhow::Result<Option<MidjourneySession>> {
+    let _guard = self.mutation.lock().await;
+    let Some(mut credential) = self.snapshot()? else {
+      return Ok(None);
     };
-    
-    let maybe_cookies= maybe_cookies
-        .map(|cookies| cookies.to_serializable());
-    
-    let maybe_user_info= maybe_user_info
-        .map(|info| info.to_serializable());
-
-    if maybe_cookies.is_none() && maybe_user_info.is_none() {
-      info!("Nothing to write to disk, skipping.");
-      return Ok(());
-      
+    if credential.cookie.maybe_relogin_required_since.is_some() {
+      return Ok(None);
     }
-
-    let state = SerializableMidjourneyState {
-      version: SERIALIZABLE_MIDJOURNEY_STATE_VERSION,
-      user_cookies: maybe_cookies,
-      user_info: maybe_user_info,
+    if token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_ID), Utc::now().timestamp()) {
+      let Some(refresh) = credential.cookie.cookies.get_cookie_value(AUTH_REFRESH).map(str::to_owned) else {
+        self.mark_relogin(&mut credential)?;
+        return Ok(None);
+      };
+      let refreshed = tokio::time::timeout(Duration::from_secs(30), refresh_token(RefreshTokenArgs { request: RefreshTokenRequest { refresh_token: &refresh, api_key: DEFAULT_FIREBASE_API_KEY }, browser: Some(credential.browser()) })).await?;
+      match refreshed {
+        Ok(response) => {
+          let origin = Url::parse(MIDJOURNEY_ORIGIN)?;
+          // Preserve the host-only Secure scope required by __Host- cookies.
+          for (name, value) in [(AUTH_ID, response.id_token), (AUTH_REFRESH, response.refresh_token)] {
+            credential.cookie.cookies.apply_set_cookie_header(&format!("{name}={value}; Path=/; Secure; HttpOnly"), &origin);
+          }
+          credential.cookie.maybe_updated_at = Some(Utc::now());
+          credential.cookie.maybe_succeeded_at = Some(Utc::now());
+          populate_account_info(&mut credential);
+          self.store(credential.clone())?;
+        },
+        Err(err) if is_revoked_refresh(&err) => {
+          self.mark_relogin(&mut credential)?;
+          return Ok(None);
+        },
+        Err(err) => {
+          credential.cookie.maybe_failed_at = Some(Utc::now());
+          self.store(credential)?;
+          return Err(err.into());
+        },
+      }
+    }
+    populate_account_info(&mut credential);
+    let cookie_header = credential.cookie.cookies.cookie_header_for_url(&Url::parse(MIDJOURNEY_ORIGIN)?).unwrap_or_default();
+    let browser = credential.browser();
+    let mut maybe_user_id = extract_midjourney_user_id_from_cookie_header(&cookie_header).or_else(|| credential.maybe_user_info.as_ref()?.maybe_user_id.as_deref().map(MidjourneyUserId::from_str));
+    if maybe_user_id.is_none() {
+      let info = tokio::time::timeout(Duration::from_secs(30), get_user_info(GetUserInfoArgs { cookie_header: &cookie_header, hostname: None, browser: Some(browser.clone()) })).await??;
+      maybe_user_id = info.user_id;
+      credential.maybe_user_info = Some(MidjourneyAccountInfo { maybe_user_id: maybe_user_id.as_ref().map(|id| id.as_str().to_string()), maybe_email: info.email, maybe_username: None });
+      self.store(credential.clone())?;
+    }
+    let Some(user_id) = maybe_user_id else {
+      bail!("Could not resolve the Midjourney account; please log in again")
     };
+    Ok(Some(MidjourneySession { credential_id: credential.id, cookie_header, user_id, browser }))
+  }
 
-    let path = self.app_data_root.credentials_dir().get_midjourney_state_path();
-    let serialized = serde_json::to_string(&state)?;
-    
-    std::fs::write(path, serialized)?;
+  fn snapshot(&self) -> anyhow::Result<Option<MidjourneyCredential>> {
+    Ok(self.credential.read().map_err(|_| anyhow!("Midjourney credential lock poisoned"))?.clone())
+  }
 
+  fn store(&self, credential: MidjourneyCredential) -> anyhow::Result<()> {
+    credential.save(&self.credential_path())?;
+    *self.credential.write().map_err(|_| anyhow!("Midjourney credential lock poisoned"))? = Some(credential);
     Ok(())
   }
+
+  fn mark_relogin(&self, credential: &mut MidjourneyCredential) -> anyhow::Result<()> {
+    let now = Utc::now();
+    credential.cookie.maybe_failed_at = Some(now);
+    credential.cookie.maybe_relogin_required_since.get_or_insert(now);
+    self.store(credential.clone())
+  }
+
+  fn credential_path(&self) -> std::path::PathBuf {
+    self.app_data_root.credentials_dir().get_midjourney_credential_path()
+  }
 }
 
-fn read_midjourney_state_from_disk(root: &AppDataRoot) -> AnyhowResult<Option<SerializableMidjourneyState>> {
-  let midjourney_state_path= root.credentials_dir().get_midjourney_state_path();
-
-  if !midjourney_state_path.exists() {
+fn load_or_migrate(root: &AppDataRoot) -> anyhow::Result<Option<MidjourneyCredential>> {
+  let path = root.credentials_dir().get_midjourney_credential_path();
+  if path.exists() {
+    let credential = MidjourneyCredential::load(&path)?;
+    credential.save(&path)?; // Persist a stable ID for hand-written files.
+    return Ok(Some(credential));
+  }
+  let legacy_path = root.credentials_dir().get_midjourney_state_path();
+  if !legacy_path.exists() {
     return Ok(None);
   }
-  
-  let contents = read_to_string(&midjourney_state_path)?
-      .trim()
-      .to_string();
-  
-  let state: SerializableMidjourneyState = serde_json::from_str(&contents)?;
-
-  Ok(Some(state))
+  let legacy: SerializableMidjourneyState = serde_json::from_str(&std::fs::read_to_string(legacy_path)?)?;
+  let Some(cookies) = legacy.user_cookies else {
+    return Ok(None);
+  };
+  let cookies = CookieStore::from_cookie_header(&cookies.to_cookie_store().to_cookie_string(), &Url::parse(MIDJOURNEY_ORIGIN)?);
+  let mut credential = MidjourneyCredential::new(cookies);
+  credential.maybe_user_info = legacy.user_info.map(|info| MidjourneyAccountInfo { maybe_user_id: info.user_id, maybe_email: info.email, maybe_username: None });
+  populate_account_info(&mut credential);
+  credential.save(&path)?;
+  Ok(Some(credential))
+}
+
+fn populate_account_info(credential: &mut MidjourneyCredential) {
+  let Some(claims) = credential.cookie.cookies.get_cookie_value(AUTH_ID).and_then(jwt_claims) else {
+    return;
+  };
+  let info = credential.maybe_user_info.get_or_insert_with(MidjourneyAccountInfo::default);
+  if let Some(id) = claims.get("midjourney_id").and_then(Value::as_str).filter(|id| !id.is_empty()) {
+    info.maybe_user_id = Some(id.to_string());
+  }
+  if let Some(email) = claims.get("email").and_then(Value::as_str) {
+    info.maybe_email = Some(email.to_string());
+  }
+}
+
+fn token_needs_refresh(maybe_token: Option<&str>, now: i64) -> bool {
+  maybe_token.and_then(jwt_claims).and_then(|claims| claims.get("exp")?.as_i64()).is_none_or(|expires| expires <= now + 120)
+}
+
+fn jwt_claims(token: &str) -> Option<Value> {
+  let payload = token.split('.').nth(1)?;
+  serde_json::from_slice(&BASE64_URL_SAFE_NO_PAD.decode(payload).ok()?).ok()
+}
+
+fn is_revoked_refresh(error: &MidjourneyError) -> bool {
+  let MidjourneyError::Api(MidjourneyApiError::InvalidRequest(body) | MidjourneyApiError::Unauthorized(body)) = error else {
+    return false;
+  };
+  ["TOKEN_EXPIRED", "INVALID_REFRESH_TOKEN", "USER_DISABLED", "USER_NOT_FOUND"].iter().any(|code| body.contains(code))
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
+  use tempfile::TempDir;
+
+  const TEST_USER_ID: &str = "00000000-0000-4000-8000-000000000001";
+
+  mod persistence {
+    use super::*;
+
+    #[tokio::test]
+    async fn migrates_legacy_cookies_and_logout_does_not_touch_artcraft() {
+      let (_dir, root) = root();
+      let legacy = root.credentials_dir().get_midjourney_state_path();
+      let artcraft = root.credentials_dir().get_storyteller_session_cookie_file_path();
+      std::fs::write(&artcraft, "existing-artcraft-session").unwrap();
+      std::fs::write(
+        &legacy,
+        serde_json::json!({
+          "version": 1,
+          "user_cookies": {"cookies": [
+            {"name": AUTH_ID, "value": token(Utc::now().timestamp() + 3600)},
+            {"name": AUTH_REFRESH, "value": "test-refresh"}
+          ]},
+          "user_info": {"user_id": TEST_USER_ID, "email": "test@example.invalid", "websocket_token": "old-socket-token"}
+        })
+        .to_string(),
+      )
+      .unwrap();
+      let manager = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
+      let session = manager.session().await.unwrap().unwrap();
+      assert_eq!(session.user_id.as_str(), TEST_USER_ID);
+      let contents = std::fs::read_to_string(root.credentials_dir().get_midjourney_credential_path()).unwrap();
+      assert!(contents.contains("midjourney_cookies"));
+      assert!(!contents.contains("old-socket-token"));
+      assert!(legacy.exists(), "migration preserves the old file until logout");
+      let reloaded = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
+      assert_eq!(reloaded.session().await.unwrap().unwrap().credential_id, session.credential_id);
+      manager.clear_credentials().await.unwrap();
+      assert!(!legacy.exists());
+      assert!(!root.credentials_dir().get_midjourney_credential_path().exists());
+      assert_eq!(std::fs::read_to_string(artcraft).unwrap(), "existing-artcraft-session");
+      let reloaded = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
+      assert!(reloaded.session().await.unwrap().is_none());
+    }
+
+    #[tokio::test]
+    async fn expired_session_without_refresh_requires_login_and_new_login_clears_it() {
+      let (_dir, root) = root();
+      let manager = MidjourneyCredentialManager::initialize_empty(&root);
+      manager.save_login(cookies(10)).await.unwrap();
+      assert!(manager.session().await.unwrap().is_none());
+      assert!(!manager.session_appears_active().unwrap());
+      let dead = MidjourneyCredential::load(&root.credentials_dir().get_midjourney_credential_path()).unwrap();
+      assert!(dead.cookie.maybe_relogin_required_since.is_some());
+      manager.save_login(cookies(Utc::now().timestamp() + 3600)).await.unwrap();
+      assert!(manager.session().await.unwrap().is_some());
+      assert!(manager.session_appears_active().unwrap());
+      assert!(manager.snapshot().unwrap().unwrap().cookie.maybe_relogin_required_since.is_none());
+    }
+
+    #[test]
+    fn cookie_attributes_and_user_agent_survive_toml_round_trip() {
+      let (_dir, root) = root();
+      let origin = Url::parse(MIDJOURNEY_ORIGIN).unwrap();
+      let mut cookies = CookieStore::empty();
+      cookies.apply_set_cookie_header("session=host; Path=/; Secure; HttpOnly; SameSite=Lax", &origin);
+      cookies.apply_set_cookie_header("session=domain; Domain=midjourney.com; Path=/api; Secure", &origin);
+      let credential = MidjourneyCredential::new(cookies);
+      let path = root.credentials_dir().get_midjourney_credential_path();
+      credential.save(&path).unwrap();
+      let loaded = MidjourneyCredential::load(&path).unwrap();
+      assert_eq!(loaded.cookie.cookies, credential.cookie.cookies);
+      assert_eq!(loaded.browser(), credential.browser());
+      assert_eq!(loaded.cookie.cookies.cookie_header_for_url(&origin).as_deref(), Some("session=host"));
+      assert!(loaded.cookie.cookies.cookie_header_for_url(&Url::parse("http://www.midjourney.com/").unwrap()).is_none());
+      assert_eq!(loaded.cookie.cookies.cookie_header_for_url(&Url::parse("https://midjourney.com/api/test").unwrap()).as_deref(), Some("session=domain"));
+      #[cfg(unix)]
+      {
+        use std::os::unix::fs::PermissionsExt;
+        assert_eq!(std::fs::metadata(path).unwrap().permissions().mode() & 0o777, 0o600);
+      }
+    }
+
+    #[test]
+    fn corrupt_modern_credential_does_not_revive_legacy_state() {
+      let (_dir, root) = root();
+      std::fs::write(root.credentials_dir().get_midjourney_credential_path(), "not valid toml").unwrap();
+      std::fs::write(root.credentials_dir().get_midjourney_state_path(), "{}").unwrap();
+      assert!(load_or_migrate(&root).is_err());
+    }
+  }
+
+  mod refresh {
+    use super::*;
+
+    #[test]
+    fn refreshes_before_expiry_but_does_not_refresh_a_fresh_token() {
+      assert!(!token_needs_refresh(Some(&token(5000)), 1000));
+      assert!(token_needs_refresh(Some(&token(1120)), 1000));
+      assert!(token_needs_refresh(Some(&token(999)), 1000));
+      assert!(token_needs_refresh(Some("invalid"), 1000));
+      assert!(token_needs_refresh(None, 1000));
+    }
+
+    #[test]
+    fn temporary_refresh_failures_do_not_force_logout() {
+      assert!(!is_revoked_refresh(&MidjourneyApiError::NetworkError("offline".into()).into()));
+      assert!(!is_revoked_refresh(&MidjourneyApiError::InvalidRequest("API_KEY_INVALID".into()).into()));
+      assert!(is_revoked_refresh(&MidjourneyApiError::InvalidRequest("INVALID_REFRESH_TOKEN".into()).into()));
+    }
+  }
+
+  fn root() -> (TempDir, AppDataRoot) {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    assert!(root.credentials_dir().path().exists());
+    (dir, root)
+  }
+
+  fn cookies(exp: i64) -> CookieStore {
+    CookieStore::from_cookie_header(&format!("{AUTH_ID}={}", token(exp)), &Url::parse(MIDJOURNEY_ORIGIN).unwrap())
+  }
+
+  fn token(exp: i64) -> String {
+    let claims = serde_json::json!({"exp": exp, "midjourney_id": TEST_USER_ID});
+    format!("header.{}.signature", BASE64_URL_SAFE_NO_PAD.encode(claims.to_string()))
+  }
 }

crates/desktop/artcraft/src/services/midjourney/state/mod.rs

@@ -1,3 +1,4 @@
 pub (super) mod serializable_midjourney_state;
 pub mod midjourney_credential_manager;
-pub mod midjourney_user_info;
+pub mod midjourney_user_info;
+pub mod midjourney_credential;

crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs

@@ -1,403 +1,269 @@
-use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
-use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
-use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
-use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
 use crate::core::state::task_database::TaskDatabase;
 use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES;
-use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
-use crate::services::midjourney::threads::events::maybe_handle_text_to_image_complete_event::maybe_handle_text_to_image_complete_event;
-use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image;
+use crate::services::midjourney::state::midjourney_credential_manager::{MidjourneyCredentialManager, MidjourneySession};
+use crate::services::midjourney::threads::upload_midjourney_batch::upload_midjourney_batch;
 use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
-use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
-use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
-use cookie_store::cookie_store::CookieStore;
 use enums::common::generation_provider::GenerationProvider;
-use enums::common::generation::common_model_type::CommonModelType;
-use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
 use errors::AnyhowResult;
-use uuid_utils::uuid::generate_random_uuid;
-use log::{error, info};
-use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
-use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
-use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest, MidjourneyJobType};
-use midjourney_client::utils::get_image_url::get_image_url;
+use midjourney_client::client::websocket::midjourney_websocket::MidjourneyWebSocket;
+use midjourney_client::client::websocket::midjourney_ws_event::MidjourneyWsEvent;
+use midjourney_client::client::websocket::open_midjourney_websocket::{open_midjourney_websocket, OpenMidjourneyWebSocketRequest};
+use midjourney_client::endpoints::imagine::{imagine, ImagineArgs, ImagineItem, ImagineRequest};
+use midjourney_client::recipes::get_user_info::{get_user_info, GetUserInfoArgs};
 use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
-use once_cell::sync::Lazy;
-use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList};
-use sqlite_tasks::queries::task::Task;
-use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
-use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
+use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs};
 use std::collections::{HashMap, HashSet};
-use std::fs::File;
-use std::io::Write;
-use std::path::PathBuf;
-use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
-use artcraft_client::endpoints::media_files::get_media_file::get_media_file;
-use artcraft_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
-use artcraft_client::endpoints::prompts::create_prompt::create_prompt;
-use artcraft_client::error::api_error::ApiError;
-use artcraft_client::error::storyteller_error::StorytellerError;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
 use tauri::AppHandle;
-use tokens::tokens::batch_generations::BatchGenerationToken;
-use url::Url;
+use tokio::sync::broadcast::{error::TryRecvError, Receiver};
+use tokio::time::timeout;
 
-/// This thread is responsible for picking up tasks that fell through the cracks of
-/// the faster websocket thread.
-pub async fn midjourney_long_polling_thread(
-  app_handle: AppHandle,
-  app_env_configs: AppEnvConfigs,
-  app_data_root: AppDataRoot,
-  task_database: TaskDatabase,
-  creds: MidjourneyCredentialManager,
-  storyteller_creds_manager: StorytellerCredentialManager,
-) -> ! {
+const RESCAN_INTERVAL: Duration = Duration::from_secs(3);
+const FEED_INTERVAL: Duration = Duration::from_secs(15);
+const SOCKET_RETRY_INTERVAL: Duration = Duration::from_secs(60);
+const REQUEST_TIMEOUT: Duration = Duration::from_secs(20);
+
+struct CompletionWorker {
+  maybe_socket: Option<SocketSession>,
+  completed: HashSet<String>,
+  feed: HashMap<String, ImagineItem>,
+  next_feed: Instant,
+  next_socket_attempt: Instant,
+}
+
+struct SocketSession {
+  credential_id: String,
+  socket: MidjourneyWebSocket,
+  events: Receiver<Arc<MidjourneyWsEvent>>,
+  subscribed: HashSet<String>,
+}
+
+/// One completion owner consumes both websocket notifications and the HTTP
+/// feed. HTTP reconciliation runs even while the socket is connected, covering
+/// missed frames and jobs enqueued before a restart or before subscription.
+pub async fn midjourney_long_polling_thread(app_handle: AppHandle, app_env_configs: AppEnvConfigs, app_data_root: AppDataRoot, task_database: TaskDatabase, credentials: MidjourneyCredentialManager, storyteller: StorytellerCredentialManager) -> ! {
+  let mut worker = CompletionWorker { maybe_socket: None, completed: HashSet::new(), feed: HashMap::new(), next_feed: Instant::now(), next_socket_attempt: Instant::now() };
   loop {
-    let res = polling_loop(
-      &app_handle,
-      &app_env_configs,
-      &app_data_root,
-      &task_database,
-      &creds,
-      &storyteller_creds_manager,
-    ).await;
-    if let Err(err) = res {
-      error!("An error occurred: {:?}", err);
+    if let Err(err) = worker.tick(Some(&app_handle), &app_env_configs, &app_data_root, &task_database, &credentials, &storyteller).await {
+      log::warn!("Midjourney completion polling failed: {}", err);
+      tokio::time::sleep(FEED_INTERVAL).await;
     }
-    // NB: Only sleep if an error occurs.
-    tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
+    tokio::time::sleep(RESCAN_INTERVAL).await;
   }
 }
 
-async fn polling_loop(
-  app_handle: &AppHandle,
-  app_env_configs: &AppEnvConfigs,
-  app_data_root: &AppDataRoot,
-  task_database: &TaskDatabase,
-  creds: &MidjourneyCredentialManager,
-  storyteller_creds_manager: &StorytellerCredentialManager,
-) -> AnyhowResult<()> {
-  loop {
-    if !creds.session_appears_active()? {
-      tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
-      continue;
+impl CompletionWorker {
+  async fn tick(&mut self, maybe_app: Option<&AppHandle>, config: &AppEnvConfigs, root: &AppDataRoot, database: &TaskDatabase, credentials: &MidjourneyCredentialManager, storyteller: &StorytellerCredentialManager) -> AnyhowResult<()> {
+    let tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs { db: database.get_connection(), provider: GenerationProvider::Midjourney, task_statuses: &TASK_DATABASE_PENDING_STATUSES }).await?.tasks;
+    if tasks.is_empty() {
+      self.close_socket();
+      self.feed.clear();
+      self.completed.clear();
+      self.next_feed = Instant::now();
+      return Ok(());
+    }
+    let Some(session) = credentials.session().await? else {
+      self.close_socket();
+      return Ok(());
+    };
+    if self.maybe_socket.as_ref().is_some_and(|socket| socket.credential_id != session.credential_id || !socket.socket.is_connected()) {
+      self.close_socket();
+    }
+    let pending: HashSet<String> = tasks.iter().filter_map(|task| task.provider_job_id.clone()).collect();
+    self.completed.retain(|job| pending.contains(job));
+    self.feed.retain(|job, _| pending.contains(job));
+    self.drain_events(&pending);
+
+    if Instant::now() >= self.next_feed {
+      self.next_feed = Instant::now() + FEED_INTERVAL;
+      match timeout(REQUEST_TIMEOUT, imagine(ImagineArgs { request: ImagineRequest { user_id: &session.user_id, page_size: None }, cookie_header: &session.cookie_header, hostname: None, browser: Some(session.browser.clone()) })).await {
+        Ok(Ok(response)) => {
+          for item in response.items {
+            if let Some(id) = item.id.clone().filter(|id| pending.contains(id)) {
+              self.feed.insert(id, item);
+            }
+          }
+        },
+        other => log::warn!("Midjourney HTTP reconciliation failed: {:?}", other),
+      }
     }
 
-    // TODO: Graceful wait, fix this long function body
-    let storyteller_creds = match storyteller_creds_manager.get_credentials()? {
-      Some(creds) => creds,
-      None => {
-        error!("No Storyteller credentials found. Cannot proceed with Midjourney polling.");
-        tokio::time::sleep(std::time::Duration::from_millis(5_000)).await;
-        continue;
+    if self.maybe_socket.is_none() && Instant::now() >= self.next_socket_attempt {
+      self.next_socket_attempt = Instant::now() + SOCKET_RETRY_INTERVAL;
+      // Resolve a fresh websocket token independently of the cached JWT user
+      // ID. Knowing the ID alone must not skip this index-page request.
+      match timeout(REQUEST_TIMEOUT, open_socket(&session)).await {
+        Ok(Ok(socket)) => self.maybe_socket = Some(socket),
+        _ => log::debug!("Midjourney websocket unavailable; HTTP polling remains active"),
       }
-    };
-
-    let cookies = creds.maybe_copy_cookie_store()?;
-
-    let cookies = match cookies {
-      Some(cookies) => cookies,
-      None => {
-        tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
-        continue;
-      }
-    };
-
-    let user_info = creds.maybe_copy_user_info()?;
-
-    let maybe_user_id = user_info
-        .map(|info| info.user_id)
-        .flatten();
-
-    let user_id = match maybe_user_id {
-      Some(user_id) => user_id,
-      None => {
-        tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
-        continue;
-      }
-    };
-
-    let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs {
-      db: task_database.get_connection(),
-      provider: GenerationProvider::Midjourney,
-      task_statuses: &TASK_DATABASE_PENDING_STATUSES,
-    }).await?;
-
-    poll_midjourney_tasks(
-      app_handle,
-      app_env_configs,
-      app_data_root,
-      task_database,
-      &cookies,
-      &user_id,
-      &storyteller_creds,
-      local_tasks,
-    ).await?;
-
-    tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
-  }
-}
-
-async fn poll_midjourney_tasks(
-  app_handle: &AppHandle,
-  app_env_configs: &AppEnvConfigs,
-  app_data_root: &AppDataRoot,
-  task_database: &TaskDatabase,
-  mj_cookies: &CookieStore,
-  mj_user_id: &MidjourneyUserId,
-  storyteller_creds: &StorytellerCredentialSet,
-  local_tasks: TaskList,
-) -> AnyhowResult<()> {
-  let local_tasks = local_tasks.tasks;
-
-  if local_tasks.is_empty() {
-    return Ok(())
-  }
-
-  // Map of Midjourney Job ID to Local Task.
-  let local_tasks_by_midjourney_job_id = local_tasks.iter()
-      .filter_map(|task| {
-        if let Some(provider_job_id) = &task.provider_job_id {
-          Some((provider_job_id.clone(), task.clone()))
-        } else {
-          None
+    }
+    if let Some(socket) = self.maybe_socket.as_mut() {
+      socket.subscribed.retain(|job| pending.contains(job));
+      for id in &pending {
+        if !socket.subscribed.contains(id) && socket.socket.subscribe_to_job(id).is_ok() {
+          socket.subscribed.insert(id.clone());
         }
-      })
-      .collect::<HashMap<String, Task>>();
-
-  let cookie_header = mj_cookies.to_cookie_string();
-
-  let midjourney_result = imagine(ImagineRequest {
-    hostname: MidjourneyHostname::Standard,
-    cookie_header,
-    user_id: mj_user_id,
-    page_size: None,
-  }).await?;
-
-  let midjourney_items = midjourney_result.items;
-
-  let midjourney_items_by_id = {
-    let mut hash = HashMap::new();
-    for item in midjourney_items.iter() {
-      if let Some(id) = &item.id {
-        hash.insert(id.to_string(), item.clone());
       }
     }
-    hash
-  };
-
-  // TODO: If we introduce another job polling mechanism, we may need to handle concurrency.
-  //  One idea might be to add a new job state that acts as an optimistic lock
-
-  let image_downloader = ImageDownloaderClient::create()?;
-
-  for (midjourney_job_id, local_task) in local_tasks_by_midjourney_job_id.iter() {
-    // TODO: Copy prompt from this.
-    let midjourney_item = match midjourney_items_by_id.get(midjourney_job_id) {
-      Some(item) => item,
-      None => continue,
+    let Some(storyteller_credentials) = storyteller.get_credentials()? else {
+      return Ok(());
     };
-
-    upload_midjourney_batch(
-      &app_handle,
-      &app_env_configs,
-      app_data_root,
-      task_database,
-      &storyteller_creds,
-      &image_downloader,
-      midjourney_job_id,
-      &local_task,
-      midjourney_item
-    ).await?;
-
-    tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
+    let downloader = ImageDownloaderClient::create(Some(session.browser))?;
+    for task in &tasks {
+      let Some(id) = task.provider_job_id.as_deref() else {
+        continue;
+      };
+      if !self.completed.contains(id) && !self.feed.contains_key(id) {
+        continue;
+      }
+      match upload_midjourney_batch(maybe_app, config, root, database, &storyteller_credentials, &downloader, id, task, self.feed.get(id)).await {
+        Ok(()) => {
+          self.completed.remove(id);
+          self.feed.remove(id);
+        },
+        Err(err) => log::warn!("Midjourney result upload for {} will be retried: {}", id, err),
+      }
+    }
+    Ok(())
   }
 
-  tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
-
-  Ok(())
-}
-
-async fn upload_midjourney_batch(
-  app_handle: &AppHandle,
-  app_env_configs: &AppEnvConfigs,
-  app_data_root: &AppDataRoot,
-  task_database: &TaskDatabase,
-  storyteller_creds: &StorytellerCredentialSet,
-  image_downloader: &ImageDownloaderClient,
-  midjourney_job_id: &String,
-  local_task: &Task,
-  midjourney_item: &ImagineItem
-) -> AnyhowResult<()> {
-  let model_type = match midjourney_item.job_type {
-    Some(MidjourneyJobType::V6Diffusion) => CommonModelType::MidjourneyV6,
-    Some(MidjourneyJobType::V6p1Diffusion) => CommonModelType::MidjourneyV6p1,
-    Some(MidjourneyJobType::V6p1RawDiffusion) => CommonModelType::MidjourneyV6p1Raw,
-    Some(MidjourneyJobType::V7Diffusion) => CommonModelType::MidjourneyV7,
-    Some(MidjourneyJobType::V7RawDiffusion) => CommonModelType::MidjourneyV7Raw,
-    Some(MidjourneyJobType::V7DraftDiffusion) => CommonModelType::MidjourneyV7Draft,
-    Some(MidjourneyJobType::V7DraftRawDiffusion) => CommonModelType::MidjourneyV7DraftRaw,
-    Some(MidjourneyJobType::Other(ref other)) => {
-      info!("Unknown Midjourney job type (for job id {}): {}", midjourney_job_id, other);
-      CommonModelType::Midjourney
-    },
-    _ => CommonModelType::Midjourney,
-  };
-
-  let request = CreatePromptRequest {
-    uuid_idempotency_token: generate_random_uuid(),
-    positive_prompt: midjourney_item.full_command.clone(),
-    negative_prompt: None,
-    model_type: Some(model_type),
-    generation_provider: Some(GenerationProvider::Midjourney),
-    maybe_generation_mode: None,
-    maybe_aspect_ratio: None,
-    maybe_resolution: None,
-    maybe_batch_count: None,
-    maybe_generate_audio: None,
-    maybe_duration_seconds: None,  };
-
-  let prompt_response = create_prompt(
-    &app_env_configs.storyteller_host,
-    Some(storyteller_creds),
-    request
-  ).await?;
-
-  info!("Created prompt: {:?}", &prompt_response.prompt_token);
-
-  // TODO: Move this from clientside to the backend.
-  //  The first upload should produce a batch token that we can reuse.
-  let batch_token = BatchGenerationToken::generate();
-
-  info!("Using synthetic batch token: {:?}", &batch_token);
-
-  let mut maybe_primary_media_file_token = None;
-
-  for index in 0..4 {
-    info!("Downloading generated Midjourney file...");
-
-    let download_path = download_midjourney_image(
-      &image_downloader,
-      midjourney_job_id,
-      index,
-      app_data_root
-    ).await?;
-
-    let mut wait_delay = 0;
-
+  fn drain_events(&mut self, pending: &HashSet<String>) {
+    let Some(socket) = self.maybe_socket.as_mut() else {
+      return;
+    };
     loop {
-      info!("Uploading to backend...");
-
-      // TODO: media_files.origin_category
-      // TODO: media_files.maybe_prompt_token
-      // TODO: media_files.maybe_generation_provider
-      // TODO: media_files.maybe_origin_model_type
-      // TODO: media_files.maybe_origin_model_token (sref?)
-      // TODO: media_files.maybe_batch_token
-      // TODO: media_files.is_user_upload
-
-      // TODO: batch_generations.token
-      // TODO: batch_generations.entity_type
-      // TODO: batch_generations.entity_token
-
-      let result = upload_image_media_file_from_file(UploadImageFromFileArgs {
-        api_host: &app_env_configs.storyteller_host,
-        maybe_creds: Some(&storyteller_creds),
-        path: &download_path,
-        is_intermediate_system_file: false,
-        maybe_prompt_token: Some(&prompt_response.prompt_token),
-        maybe_batch_token: Some(&batch_token),
-        maybe_generation_provider: Some(GenerationProvider::Midjourney),
-      }).await;
-
-      match result {
-        Ok(result) => {
-          info!("Successfully uploaded to backend: {:?}", result.media_file_token);
-          if maybe_primary_media_file_token.is_none() {
-            maybe_primary_media_file_token = Some(result.media_file_token);
+      match socket.events.try_recv() {
+        Ok(event) => {
+          if let MidjourneyWsEvent::Completed { job_id, .. } = event.as_ref() {
+            if pending.contains(job_id) {
+              self.completed.insert(job_id.clone());
+            }
           }
-          break;
         },
-        Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => {
-          error!("Too many requests, retrying upload after delay...");
-          // If we hit a rate limit, we can retry after a short delay.
-          wait_delay += 10;
-          if wait_delay > 60 {
-            wait_delay = 60;
-          }
-          tokio::time::sleep(std::time::Duration::from_secs(wait_delay)).await;
-          continue; // Retry the upload.
-        }
-        Err(err) => {
-          error!("Failed to upload to backend: {:?}", err);
-          return Err(err.into())
-        },
-      }
-    } // End loop
-
-    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
-  }
-
-  let mut maybe_cdn_url = None;
-  let mut maybe_thumbnail_url_template = None;
-
-  if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() {
-    info!("Looking up file to grab CDN and thumbnail URLs: {:?} ...", media_file_token);
-
-    let lookup_result = get_media_file(
-      &app_env_configs.storyteller_host,
-      media_file_token,
-    ).await;
-    match lookup_result {
-      Ok(response) => {
-        maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string());
-        maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links)
-            .map(|s| s.to_string());
-      }
-      Err(err) => {
-        error!("Failed to look up media file after upload: {:?} (failing open)", err);
+        Err(TryRecvError::Lagged(_)) => self.next_feed = Instant::now(),
+        Err(TryRecvError::Empty | TryRecvError::Closed) => break,
       }
     }
   }
 
-  let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
-    db: task_database.get_connection(),
-    task_id: &local_task.id,
-    maybe_batch_token: Some(&batch_token),
-    maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(),
-    maybe_primary_media_file_class: Some(TaskMediaFileClass::Image),
-    maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(),
-    maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(),
-  }).await?;
-
-  if !updated {
-    return Ok(()); // If anything breaks with queries, don't spam events.
+  fn close_socket(&mut self) {
+    if let Some(socket) = self.maybe_socket.take() {
+      socket.socket.close();
+    }
   }
-
-  let event = GenerationCompleteEvent {
-    //media_file_token: result.media_file_token,
-    action: Some(GenerationAction::GenerateImage),
-    service: GenerationServiceProvider::Midjourney,
-    model: None,
-  };
-
-  if let Err(err) = event.send(&app_handle) {
-    error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
-  }
-
-  let result = maybe_handle_text_to_image_complete_event(
-    app_handle,
-    app_env_configs,
-    Some(storyteller_creds),
-    local_task,
-    &batch_token,
-  ).await;
-
-  if let Err(err) = result {
-    error!("Failed to send text-to-image complete event: {:?}", err);
-  }
-
-  Ok(())
 }
 
+async fn open_socket(session: &MidjourneySession) -> AnyhowResult<SocketSession> {
+  let info = get_user_info(GetUserInfoArgs { cookie_header: &session.cookie_header, hostname: None, browser: Some(session.browser.clone()) }).await?;
+  let token = info.websocket_token.ok_or_else(|| anyhow::anyhow!("No Midjourney websocket token"))?;
+  let socket = open_midjourney_websocket(OpenMidjourneyWebSocketRequest { websocket_token: &token, user_id: session.user_id.clone(), hostname: None, browser: Some(session.browser.clone()) }).await?;
+  let events = socket.events();
+  Ok(SocketSession { credential_id: session.credential_id.clone(), socket, events, subscribed: HashSet::new() })
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use crate::core::commands::generate::generate_image::providers::midjourney::handle_midjourney::handle_midjourney;
+  use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
+  use crate::services::midjourney::state::midjourney_credential::MidjourneyCredential;
+  use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
+  use artcraft_client::endpoints::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
+  use artcraft_client::utils::api_host::ApiHost;
+  use enums::tauri::tasks::task_status::TaskStatus;
+  use sqlite_tasks::queries::list_tasks_for_frontend::list_tasks_for_frontend;
+
+  /// Opt-in live test: submits ONE paid Midjourney job and uploads its four
+  /// results. An explicit scratch data root is required and reused on retry,
+  /// so rerunning after a polling/upload failure never submits another job.
+  #[ignore = "requires explicit authorization, live account cookies, and network access"]
+  #[tokio::test]
+  async fn midjourney_live_full_flow() -> AnyhowResult<()> {
+    let root = AppDataRoot::create_existing(std::env::var("ARTCRAFT_MIDJOURNEY_TEST_ROOT")?)?;
+    let target = root.credentials_dir().get_midjourney_credential_path();
+    if !target.exists() {
+      let source = std::env::var("ARTCRAFT_MIDJOURNEY_TEST_COOKIES")?;
+      MidjourneyCredential::load(std::path::Path::new(&source))?.save(&target)?;
+    }
+    let credentials = MidjourneyCredentialManager::initialize_from_disk_infallible(&root);
+    let header = std::fs::read_to_string(std::env::var("ARTCRAFT_TEST_BACKEND_COOKIES")?)?;
+    let storyteller_credentials = StorytellerCredentialSet::parse_multi_cookie_header(header.trim())?.ok_or_else(|| anyhow::anyhow!("No ArtCraft backend cookies"))?;
+    let storyteller = StorytellerCredentialManager::initialize_empty(&root);
+    storyteller.set_credentials(&storyteller_credentials)?;
+    let config = AppEnvConfigs { storyteller_host: ApiHost::Storyteller };
+    let database = TaskDatabase::connect(&root).await?;
+    let session = credentials.session().await?.ok_or_else(|| anyhow::anyhow!("Midjourney login required"))?;
+    println!("Midjourney credential loaded and refreshed; source credential unchanged");
+    let maybe_socket = match timeout(REQUEST_TIMEOUT, open_socket(&session)).await {
+      Ok(Ok(socket)) => {
+        println!("Midjourney websocket handshake succeeded");
+        Some(socket)
+      },
+      Ok(Err(err)) => {
+        println!("Midjourney websocket unavailable: {}; testing HTTP fallback", err);
+        None
+      },
+      Err(_) => {
+        println!("Midjourney websocket connection timed out; testing HTTP fallback");
+        None
+      },
+    };
+    let mut maybe_events = maybe_socket.as_ref().map(|socket| socket.socket.events());
+    if list_tasks_for_frontend(database.get_connection()).await?.tasks.is_empty() {
+      let request: TauriGenerateImageRequest = serde_json::from_value(serde_json::json!({
+        "model": "midjourney_8", "provider": "midjourney", "batch_size": 4,
+        "prompt": "A small wooden sailboat on a calm blue lake, soft morning light, watercolor illustration",
+        "aspect_ratio": "wide_sixteen_by_nine",
+        "frontend_caller": "text_to_image", "frontend_subscriber_id": "midjourney-live-test",
+      }))?;
+      let success = handle_midjourney(&request, &config, &credentials, &storyteller).await.map_err(|err| anyhow::anyhow!("Midjourney enqueue failed: {:?}", err))?;
+      let task_id = success.insert_into_task_database_with_frontend_payload(&database, request.frontend_caller, request.frontend_subscriber_id.as_deref(), None).await?;
+      println!("Enqueued Midjourney job {} as task {}", success.provider_job_id.as_deref().unwrap_or_default(), task_id.as_str());
+    } else {
+      println!("Resuming existing live-test task; no new generation submitted");
+    }
+    let mut worker = CompletionWorker { maybe_socket, completed: HashSet::new(), feed: HashMap::new(), next_feed: Instant::now(), next_socket_attempt: Instant::now() + SOCKET_RETRY_INTERVAL };
+    let deadline = Instant::now() + Duration::from_secs(360);
+    let mut websocket_completed = false;
+    loop {
+      worker.tick(None, &config, &root, &database, &credentials, &storyteller).await?;
+      if let Some(events) = maybe_events.as_mut() {
+        while let Ok(event) = events.try_recv() {
+          if matches!(event.as_ref(), MidjourneyWsEvent::Completed { .. }) {
+            websocket_completed = true;
+          }
+        }
+      }
+      let tasks = list_tasks_for_frontend(database.get_connection()).await?.tasks;
+      let task = tasks.first().ok_or_else(|| anyhow::anyhow!("No live-test task"))?;
+      if task.status == TaskStatus::CompleteSuccess {
+        let batch = task.on_complete_batch_token.as_ref().ok_or_else(|| anyhow::anyhow!("Missing completed batch"))?;
+        let files = list_batch_generated_redux_media_files(&config.storyteller_host, Some(&storyteller_credentials), batch).await?;
+        assert_eq!(files.media_files.len(), 4, "all four results must be uploaded");
+        assert!(task.on_complete_primary_media_file_cdn_url.is_some());
+        assert_eq!(task.provider, Some(GenerationProvider::Midjourney));
+        let job_id = task.provider_job_id.as_deref().unwrap();
+        // Independently exercise the recovery feed even when websocket won.
+        let feed = imagine(ImagineArgs { request: ImagineRequest { user_id: &session.user_id, page_size: None }, cookie_header: &session.cookie_header, hostname: None, browser: Some(session.browser.clone()) }).await?;
+        assert!(feed.items.iter().any(|item| item.id.as_deref() == Some(job_id)), "HTTP feed must find the completed job");
+        println!("Live flow succeeded: websocket_completed={}, HTTP reconciliation verified, batch={}, uploaded_files={}", websocket_completed, batch.as_str(), files.media_files.len());
+        for file in files.media_files {
+          let prompt = file.maybe_prompt_token.as_ref().ok_or_else(|| anyhow::anyhow!("Uploaded image has no prompt attribution"))?;
+          let prompt: serde_json::Value = reqwest::Client::new().get(format!("{}/v1/prompts/{}", config.storyteller_host.to_api_hostname_and_scheme(), prompt.as_str())).header("Cookie", storyteller_credentials.maybe_as_cookie_header().unwrap_or_default()).send().await?.error_for_status()?.json().await?;
+          assert_eq!(prompt["prompt"]["maybe_generation_provider"], "midjourney");
+          assert_eq!(prompt["prompt"]["maybe_model_type"], "midjourney_8");
+          println!("Uploaded result: {} {}; backend prompt provider=midjourney model=midjourney_8", file.token.as_str(), file.media_links.cdn_url);
+        }
+        worker.close_socket();
+        return Ok(());
+      }
+      if Instant::now() >= deadline {
+        anyhow::bail!("Timed out waiting for Midjourney completion; rerun to resume the same task")
+      }
+      tokio::time::sleep(RESCAN_INTERVAL).await;
+    }
+  }
+}

crates/desktop/artcraft/src/services/midjourney/threads/mod.rs

@@ -1,2 +1,3 @@
 pub mod midjourney_long_polling_thread;
-mod events;
+mod events;
+mod upload_midjourney_batch;

crates/desktop/artcraft/src/services/midjourney/threads/upload_midjourney_batch.rs

@@ -0,0 +1,202 @@
+use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
+use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
+use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
+use crate::core::state::data_dir::app_data_root::AppDataRoot;
+use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
+use crate::core::state::task_database::TaskDatabase;
+use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint};
+use crate::services::midjourney::threads::events::maybe_handle_text_to_image_complete_event::maybe_handle_text_to_image_complete_event;
+use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image;
+use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
+use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
+use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
+use artcraft_client::endpoints::media_files::get_media_file::get_media_file;
+use artcraft_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file_with_idempotency_token, UploadImageFromFileArgs};
+use artcraft_client::endpoints::prompts::create_prompt::create_prompt;
+use enums::common::generation::common_model_type::CommonModelType;
+use enums::common::generation_provider::GenerationProvider;
+use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass;
+use enums::tauri::tasks::task_model_type::TaskModelType;
+use errors::AnyhowResult;
+use log::{error, info};
+use midjourney_client::endpoints::imagine::{ImagineItem, MidjourneyJobType};
+use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
+use serde::{Deserialize, Serialize};
+use sqlite_tasks::queries::task::Task;
+use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs};
+use std::io::Write;
+use std::path::PathBuf;
+use std::time::Duration;
+use tauri::AppHandle;
+use tempfile::NamedTempFile;
+use tokens::tokens::batch_generations::BatchGenerationToken;
+use tokens::tokens::media_files::MediaFileToken;
+use tokens::tokens::prompts::PromptToken;
+use uuid_utils::uuid::generate_random_uuid;
+
+const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
+
+#[derive(Serialize, Deserialize)]
+struct MidjourneyUploadCheckpoint {
+  batch_token: BatchGenerationToken,
+  prompt_idempotency_token: String,
+  maybe_prompt_token: Option<PromptToken>,
+  #[serde(default = "upload_idempotency_tokens")]
+  upload_idempotency_tokens: [String; 4],
+  uploaded: Vec<MediaFileToken>,
+}
+
+pub(super) async fn upload_midjourney_batch(maybe_app_handle: Option<&AppHandle>, app_env_configs: &AppEnvConfigs, app_data_root: &AppDataRoot, task_database: &TaskDatabase, storyteller_creds: &StorytellerCredentialSet, image_downloader: &ImageDownloaderClient, midjourney_job_id: &str, local_task: &Task, maybe_item: Option<&ImagineItem>) -> AnyhowResult<()> {
+  let mut checkpoint = MidjourneyUploadCheckpoint::load_or_create(app_data_root, local_task)?;
+  let model_type = match local_task.model_type {
+    Some(TaskModelType::Midjourney7) => CommonModelType::Midjourney7,
+    Some(TaskModelType::Midjourney7Niji) => CommonModelType::Midjourney7Niji,
+    Some(TaskModelType::Midjourney8) => CommonModelType::Midjourney8,
+    _ => model_type_from_feed(maybe_item),
+  };
+  if checkpoint.maybe_prompt_token.is_none() {
+    let request = CreatePromptRequest { uuid_idempotency_token: checkpoint.prompt_idempotency_token.clone(), positive_prompt: maybe_item.and_then(|item| item.full_command.clone()), negative_prompt: None, model_type: Some(model_type), generation_provider: Some(GenerationProvider::Midjourney), maybe_generation_mode: None, maybe_aspect_ratio: None, maybe_resolution: None, maybe_batch_count: Some(4), maybe_generate_audio: None, maybe_duration_seconds: None };
+    let prompt = tokio::time::timeout(REQUEST_TIMEOUT, create_prompt(&app_env_configs.storyteller_host, Some(storyteller_creds), request)).await??;
+    checkpoint.maybe_prompt_token = Some(prompt.prompt_token);
+    checkpoint.save(app_data_root, local_task)?;
+  }
+  for index in checkpoint.uploaded.len()..4 {
+    let download_path = tokio::time::timeout(REQUEST_TIMEOUT, download_midjourney_image(image_downloader, midjourney_job_id, index as u8, app_data_root)).await??;
+    let result = tokio::time::timeout(REQUEST_TIMEOUT, upload_image_media_file_from_file_with_idempotency_token(UploadImageFromFileArgs { api_host: &app_env_configs.storyteller_host, maybe_creds: Some(storyteller_creds), path: &download_path, is_intermediate_system_file: false, maybe_prompt_token: checkpoint.maybe_prompt_token.as_ref(), maybe_batch_token: Some(&checkpoint.batch_token), maybe_generation_provider: Some(GenerationProvider::Midjourney) }, &checkpoint.upload_idempotency_tokens[index])).await??;
+    checkpoint.uploaded.push(result.media_file_token);
+    // Resume at the next image if a later upload fails or the app restarts.
+    checkpoint.save(app_data_root, local_task)?;
+  }
+  let batch_token = checkpoint.batch_token;
+  let maybe_primary_media_file_token = checkpoint.uploaded.first().cloned();
+
+  let mut maybe_cdn_url = None;
+  let mut maybe_thumbnail_url_template = None;
+
+  if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() {
+    info!("Looking up file to grab CDN and thumbnail URLs: {:?} ...", media_file_token);
+
+    let lookup_result = tokio::time::timeout(REQUEST_TIMEOUT, get_media_file(&app_env_configs.storyteller_host, media_file_token)).await;
+    match lookup_result {
+      Ok(Ok(response)) => {
+        maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string());
+        maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links).map(|s| s.to_string());
+      },
+      other => {
+        error!("Failed to look up media file after upload: {:?} (failing open)", other);
+      },
+    }
+  }
+
+  if let Some(app) = maybe_app_handle {
+    auto_download_task(app, local_task, Some(&batch_token), maybe_cdn_url.as_deref(), Some(checkpoint.uploaded.len())).await?;
+  }
+
+  let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { db: task_database.get_connection(), task_id: &local_task.id, maybe_batch_token: Some(&batch_token), maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), maybe_primary_media_file_class: Some(TaskMediaFileClass::Image), maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(), maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref() }).await?;
+
+  if !updated {
+    return Ok(()); // If anything breaks with queries, don't spam events.
+  }
+
+  if let Some(app_handle) = maybe_app_handle {
+    clear_auto_download_checkpoint(app_handle, local_task);
+    let event = GenerationCompleteEvent {
+      //media_file_token: result.media_file_token,
+      action: Some(GenerationAction::GenerateImage),
+      service: GenerationServiceProvider::Midjourney,
+      model: None,
+    };
+
+    if let Err(err) = event.send(&app_handle) {
+      error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
+    }
+
+    let result = maybe_handle_text_to_image_complete_event(app_handle, app_env_configs, Some(storyteller_creds), local_task, &batch_token).await;
+
+    if let Err(err) = result {
+      error!("Failed to send text-to-image complete event: {:?}", err);
+    }
+  }
+
+  if let Err(err) = std::fs::remove_file(MidjourneyUploadCheckpoint::path(app_data_root, local_task)) {
+    log::warn!("Could not remove completed Midjourney upload checkpoint: {}", err);
+  }
+  Ok(())
+}
+
+impl MidjourneyUploadCheckpoint {
+  fn load_or_create(root: &AppDataRoot, task: &Task) -> AnyhowResult<Self> {
+    let path = Self::path(root, task);
+    if path.exists() {
+      let checkpoint: Self = serde_json::from_str(&std::fs::read_to_string(path)?)?;
+      checkpoint.save(root, task)?; // Backfill stable tokens for older checkpoints.
+      return Ok(checkpoint);
+    }
+    let checkpoint = Self { batch_token: BatchGenerationToken::generate(), prompt_idempotency_token: generate_random_uuid(), maybe_prompt_token: task.prompt_token.as_deref().map(PromptToken::new_from_str), upload_idempotency_tokens: upload_idempotency_tokens(), uploaded: Vec::new() };
+    checkpoint.save(root, task)?;
+    Ok(checkpoint)
+  }
+
+  fn save(&self, root: &AppDataRoot, task: &Task) -> AnyhowResult<()> {
+    let mut file = NamedTempFile::new_in(root.state_dir().path())?;
+    file.write_all(&serde_json::to_vec(self)?)?;
+    file.as_file().sync_all()?;
+    file.persist(Self::path(root, task))?;
+    Ok(())
+  }
+
+  fn path(root: &AppDataRoot, task: &Task) -> PathBuf {
+    root.state_dir().path().join(format!("midjourney_upload_{}.json", task.id.as_str()))
+  }
+}
+
+fn upload_idempotency_tokens() -> [String; 4] {
+  std::array::from_fn(|_| generate_random_uuid())
+}
+
+fn model_type_from_feed(maybe_item: Option<&ImagineItem>) -> CommonModelType {
+  match maybe_item.and_then(|item| item.job_type.as_ref()) {
+    Some(MidjourneyJobType::V6Diffusion) => CommonModelType::MidjourneyV6,
+    Some(MidjourneyJobType::V6p1Diffusion) => CommonModelType::MidjourneyV6p1,
+    Some(MidjourneyJobType::V6p1RawDiffusion) => CommonModelType::MidjourneyV6p1Raw,
+    Some(MidjourneyJobType::V7Diffusion) => CommonModelType::MidjourneyV7,
+    Some(MidjourneyJobType::V7RawDiffusion) => CommonModelType::MidjourneyV7Raw,
+    Some(MidjourneyJobType::V7DraftDiffusion) => CommonModelType::MidjourneyV7Draft,
+    Some(MidjourneyJobType::V7DraftRawDiffusion) => CommonModelType::MidjourneyV7DraftRaw,
+    _ => CommonModelType::Midjourney,
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use enums::tauri::tasks::{task_status::TaskStatus, task_type::TaskType};
+  use tokens::tokens::sqlite::tasks::TaskId;
+
+  #[test]
+  fn partial_upload_resumes_with_original_prompt_batch_and_completed_images() {
+    let dir = tempfile::tempdir().unwrap();
+    let root = AppDataRoot::create_existing(dir.path()).unwrap();
+    let task = task();
+    let mut checkpoint = MidjourneyUploadCheckpoint::load_or_create(&root, &task).unwrap();
+    let batch = checkpoint.batch_token.clone();
+    let idempotency = checkpoint.prompt_idempotency_token.clone();
+    let upload_tokens = checkpoint.upload_idempotency_tokens.clone();
+    checkpoint.uploaded.push(MediaFileToken::new_from_str("mf_first"));
+    checkpoint.uploaded.push(MediaFileToken::new_from_str("mf_second"));
+    checkpoint.save(&root, &task).unwrap();
+    let loaded = MidjourneyUploadCheckpoint::load_or_create(&root, &task).unwrap();
+    assert_eq!(loaded.batch_token, batch);
+    assert_eq!(loaded.upload_idempotency_tokens, upload_tokens);
+    assert_eq!(loaded.prompt_idempotency_token, idempotency);
+    assert_eq!(loaded.maybe_prompt_token.unwrap().as_str(), "prompt_test");
+    assert_eq!(loaded.uploaded.len(), 2);
+    assert_eq!(loaded.uploaded[0].as_str(), "mf_first");
+    assert_eq!((loaded.uploaded.len()..4).collect::<Vec<_>>(), vec![2, 3]);
+  }
+
+  fn task() -> Task {
+    Task { id: TaskId::generate(), status: TaskStatus::Pending, task_type: TaskType::ImageGeneration, model_type: Some(TaskModelType::Midjourney8), provider: GenerationProvider::Midjourney, provider_job_id: Some("test-job".to_string()), queue_status_url: None, queue_response_url: None, prompt_token: Some("prompt_test".to_string()), frontend_caller: None, frontend_subscriber_id: None, frontend_subscriber_payload: None }
+  }
+}

crates/desktop/artcraft/src/services/midjourney/utils/extract_midjourney_user_id_from_cookies.rs

@@ -0,0 +1,62 @@
+use base64::prelude::BASE64_URL_SAFE_NO_PAD;
+use base64::Engine;
+use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
+use serde_json::Value;
+
+/// The Midjourney auth id-token cookie. Its (unverified) JWT payload carries
+/// the Midjourney user id under `midjourney_id`.
+const AUTH_ID_TOKEN_COOKIE: &str = "__Host-Midjourney.AuthUserTokenV3_i";
+
+/// The JWT claim holding the Midjourney user id (the `singleplayer_{...}`
+/// channel id is built from it).
+const MIDJOURNEY_ID_CLAIM: &str = "midjourney_id";
+
+/// Recover the Midjourney user id from the auth cookie's JWT, WITHOUT hitting
+/// the Cloudflare-gated index page. `cookie_header` is a raw `Cookie:` header
+/// string (`"a=1; b=2"`).
+pub fn extract_midjourney_user_id_from_cookie_header(cookie_header: &str) -> Option<MidjourneyUserId> {
+  let token = cookie_value(cookie_header, AUTH_ID_TOKEN_COOKIE)?;
+  let claims = decode_jwt_payload(token)?;
+  let id = claims.get(MIDJOURNEY_ID_CLAIM).and_then(Value::as_str)?;
+  if id.is_empty() {
+    return None;
+  }
+  Some(MidjourneyUserId::from_str(id))
+}
+
+fn cookie_value<'a>(cookie_header: &'a str, name: &str) -> Option<&'a str> {
+  cookie_header.split(';').filter_map(|pair| pair.trim().split_once('=')).find(|(cookie_name, _)| *cookie_name == name).map(|(_, value)| value)
+}
+
+/// Decode the (unverified) payload of a `header.payload.signature` JWT.
+fn decode_jwt_payload(token: &str) -> Option<Value> {
+  let mut parts = token.split('.');
+  let _header = parts.next()?;
+  let payload = parts.next()?;
+  let _signature = parts.next()?;
+  if parts.next().is_some() {
+    return None;
+  }
+  let bytes = BASE64_URL_SAFE_NO_PAD.decode(payload).ok()?;
+  serde_json::from_slice(&bytes).ok()
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  // Payload: {"midjourney_id":"26c9d38d-1f71-49c7-a356-b29dac58b54c","name":"u1"}
+  const JWT: &str = "eyJhbGciOiJIUzI1NiJ9.eyJtaWRqb3VybmV5X2lkIjoiMjZjOWQzOGQtMWY3MS00OWM3LWEzNTYtYjI5ZGFjNThiNTRjIiwibmFtZSI6InUxIn0.sig";
+
+  #[test]
+  fn extracts_user_id_from_auth_cookie() {
+    let header = format!("_fbp=abc; {}={}; other=x", AUTH_ID_TOKEN_COOKIE, JWT);
+    let user_id = extract_midjourney_user_id_from_cookie_header(&header).unwrap();
+    assert_eq!(user_id.as_str(), "26c9d38d-1f71-49c7-a356-b29dac58b54c");
+  }
+
+  #[test]
+  fn returns_none_without_auth_cookie() {
+    assert!(extract_midjourney_user_id_from_cookie_header("_fbp=abc; other=x").is_none());
+  }
+}

crates/desktop/artcraft/src/services/midjourney/utils/midjourney_browser_profile.rs

@@ -0,0 +1,21 @@
+use browser_emulation::browser_profile::BrowserProfile;
+
+/// The User-Agent used for BOTH the Midjourney login webview and every
+/// first-party Midjourney HTTP/websocket call.
+///
+/// It must be identical in both places: Cloudflare's `cf_clearance` cookie
+/// (captured by the webview) is bound to the exact User-Agent, so the wreq
+/// calls that replay those cookies have to present the same string. A
+/// mainstream desktop Safari UA (rather than the default WKWebView UA) also
+/// makes Google's embedded-webview risk detection less likely to force a
+/// passkey step-up during sign-in.
+pub const MIDJOURNEY_USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.6 Safari/605.1.15";
+
+/// The browser profile used for ALL first-party Midjourney HTTP/websocket
+/// traffic (submit, imagine, index page, image downloads, websocket upgrade).
+///
+/// Safari fingerprint + the exact [`MIDJOURNEY_USER_AGENT`], to stay consistent
+/// with the macOS WKWebView login window that captured the cookies.
+pub fn midjourney_browser_profile() -> BrowserProfile {
+  BrowserProfile::safari_macos_with_user_agent(MIDJOURNEY_USER_AGENT)
+}

crates/desktop/artcraft/src/services/midjourney/utils/mod.rs

@@ -1 +1,3 @@
-pub mod download_midjourney_image;
+pub mod download_midjourney_image;
+pub mod midjourney_browser_profile;
+pub mod extract_midjourney_user_id_from_cookies;

crates/desktop/artcraft/src/services/midjourney/windows/extract_midjourney_webview_cookies.rs

@@ -1,46 +1,17 @@
-use cookie_store::cookie_store::CookieStore;
+use cookie_store_wrapper::cookie_store::CookieStore;
 use errors::AnyhowResult;
-use once_cell::sync::Lazy;
-use reqwest::Url;
-use std::collections::HashSet;
-use tauri::webview::Cookie;
 use tauri::WebviewWindow;
-
-static WWW_COOKIE_URL: Lazy<Url> = Lazy::new(|| {
-  Url::parse("https://www.midjourney.com").expect("URL should parse")
-});
-
-static ROOT_COOKIE_URL: Lazy<Url> = Lazy::new(|| {
-  Url::parse("https://midjourney.com").expect("URL should parse")
-});
+use url::Url;
 
 pub fn extract_midjourney_webview_cookies(webview: &WebviewWindow) -> AnyhowResult<CookieStore> {
-  let mut cookie_store = CookieStore::empty();
-  let cookies = get_all_midjourney_cookies(webview)?;
-  for cookie in cookies.iter() {
-    cookie_store.add_cookie_name_and_value(
-      cookie.name().to_string(),
-      cookie.value().to_string(),
-    );
-  }
-  Ok(cookie_store)
-}
-
-fn get_all_midjourney_cookies(webview: &WebviewWindow) -> AnyhowResult<Vec<Cookie>> {
-  // NB: WWW domain contains the auth/session cookies.
-  // Root domain contains some other cookies, such as Cloudflare.
-  let www_cookies = webview.cookies_for_url(WWW_COOKIE_URL.clone())?;
-  let root_cookies = webview.cookies_for_url(ROOT_COOKIE_URL.clone())?;
-
-  let mut all_cookies = www_cookies;
-  let mut cookie_names = HashSet::new();
-
-  for cookie in root_cookies.iter() {
-    if !cookie_names.contains(cookie.name()) {
-      cookie_names.insert(cookie.name().to_string());
-      all_cookies.push(cookie.clone());
+  let mut store = CookieStore::empty();
+  for origin in ["https://www.midjourney.com/", "https://midjourney.com/"] {
+    let origin = Url::parse(origin)?;
+    for cookie in webview.cookies_for_url(origin.clone())? {
+      // Tauri's cookie serialization retains path, domain, Secure, HttpOnly,
+      // SameSite, and expiry. The jar deduplicates by domain/path/name.
+      store.apply_set_cookie_header(&cookie.to_string(), &origin);
     }
   }
-
-  Ok(all_cookies)
+  Ok(store)
 }

crates/desktop/artcraft/src/services/midjourney/windows/midjourney_login_window_thread.rs

@@ -1,185 +1,47 @@
 use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
 use crate::core::events::functional_events::refresh_account_state_event::RefreshAccountStateEvent;
-use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
-use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
-use crate::core::events::sendable_event_trait::SendableEvent;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::utils::window::get_webview_window_hostname::get_webview_window_hostname;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
-use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
 use crate::services::midjourney::windows::extract_midjourney_webview_cookies::extract_midjourney_webview_cookies;
 use crate::services::midjourney::windows::open_midjourney_login_window::MIDJOURNEY_LOGIN_WINDOW_NAME;
-use crate::services::sora::events::sora_login_success_event::SoraLoginSuccessEvent;
-use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
-use crate::services::sora::windows::sora_login_window::extract_sora_webview_cookies::extract_sora_webview_cookies;
-use anyhow::anyhow;
-use cookie_store::cookie_store::CookieStore;
 use enums::common::generation_provider::GenerationProvider;
 use errors::AnyhowResult;
-use log::{error, info};
-use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
 use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
-use midjourney_client::recipes::get_user_info::{get_user_info, GetUserInfoRequest};
-use openai_sora_client::creds::sora_credential_set::SoraCredentialSet;
-use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
-use openai_sora_client::utils::has_session_cookie::{has_session_cookie, SessionCookiePresence};
+use std::time::Duration;
 use tauri::{AppHandle, Manager, WebviewWindow};
 
-pub async fn midjourney_login_window_thread(
-  app: AppHandle,
-  app_data_root: AppDataRoot,
-  mj_creds_manager: MidjourneyCredentialManager,
-) {
-  let mut visited_login = false;
-
+pub async fn midjourney_login_window_thread(app: AppHandle, _app_data_root: AppDataRoot, credentials: MidjourneyCredentialManager) {
   loop {
-    let login_webview_window = match app.get_webview_window(MIDJOURNEY_LOGIN_WINDOW_NAME) {
-      Some(webview) => webview,
-      None => {
-        info!("Exit midjourney login thread.");
-        return; // NB: Only exit if we don't have the webview.
-      }
+    let Some(window) = app.get_webview_window(MIDJOURNEY_LOGIN_WINDOW_NAME) else {
+      return;
     };
-
-    let result = check_login_window(
-      &app,
-      &login_webview_window,
-      &app_data_root,
-      &mj_creds_manager,
-      &mut visited_login,
-    ).await;
-
-    match result {
-      Err(err) => {
-        error!("Error checking login window: {:?}", err);
-      }
-      Ok(false) => {} // Continue iteration and try again...
+    match check_login_window(&app, &window, &credentials).await {
       Ok(true) => {
-        info!("Successfully saved cookies from login window. Closing.");
-        if let Err(err) = login_webview_window.close() {
-          error!("Error closing login window: {:?}", err);
+        if let Err(err) = window.close() {
+          log::warn!("Could not close Midjourney login window: {}", err);
         }
         return;
-      }
+      },
+      Ok(false) => {},
+      Err(err) => log::warn!("Could not capture Midjourney login: {}", err),
     }
-
-    tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
+    tokio::time::sleep(Duration::from_secs(2)).await;
   }
 }
 
-/// Returns true if we can exit.
-async fn check_login_window(
-  app_handle: &AppHandle,
-  webview_window: &WebviewWindow,
-  app_data_root: &AppDataRoot,
-  mj_creds_manager: &MidjourneyCredentialManager,
-  visited_login: &mut bool,
-) -> AnyhowResult<bool> {
-
-  /* Login flow looks like this:
-
-  1. Start: https://www.midjourney.com/auth/signin
-  2. SSO / Google Login (etc): https://accounts.google.com/v3/signin/challenge/pwd?[query]
-  3. Done / Landing: https://www.midjourney.com/...
-   */
-
-  let hostname = get_webview_window_hostname(webview_window)?;
-
-  let mut maybe_at_destination = false;
-
-  match hostname.as_str() {
-    "www.midjourney.com" |
-    "midjourney.com"
-    => {
-      maybe_at_destination = true;
-    }
-    // chatgpt.com/auth is also an auth domain
-    "auth.openai.com" |
-    "accounts.google.com" |
-    "accounts.youtube.com" |
-    "login.live.com" |
-    "appleid.apple.com"
-    => {
-      // NB: We're in auth flow.
-      info!("Midjourney webview is in auth flow; hostname `{}`.", hostname);
-      *visited_login = true;
-      return Ok(false)
-    }
-    _ => {}, // We just don't know...
-  }
-
-  let cookie_store = extract_midjourney_webview_cookies(webview_window)?;
-
-  let maybe_has_auth_cookies = cookie_store_has_auth_cookies(&cookie_store);
-  let maybe_has_enough_cookies = cookie_store.len() > 6;
-  let maybe_completed_login_cycle = *visited_login && maybe_at_destination;
-
-  // Misc cookies without login cookies are ~1055 length
-  // AUTH_I is ~1500 length
-  // AUTH_R is ~500 length
-  let maybe_has_big_enough_cookie = cookie_store.calculate_approx_cookie_character_length() > 2100;
-
-  let mut heuristic_count = 0;
-
-  if maybe_has_auth_cookies {
-    heuristic_count += 1;
-  }
-  if maybe_has_enough_cookies {
-    heuristic_count += 1;
-  }
-  if maybe_completed_login_cycle {
-    heuristic_count += 1;
-  }
-  if maybe_has_big_enough_cookie {
-    heuristic_count += 1;
-  }
-
-  if heuristic_count < 3 {
+async fn check_login_window(app: &AppHandle, window: &WebviewWindow, credentials: &MidjourneyCredentialManager) -> AnyhowResult<bool> {
+  let hostname = get_webview_window_hostname(window)?;
+  if !matches!(hostname.as_str(), "www.midjourney.com" | "midjourney.com") {
     return Ok(false);
   }
-
-  info!("Heuristic count is ({}); we're likely done.", heuristic_count);
-
-  info!("Current cookies (len {}): {:?}", cookie_store.len(), cookie_store.to_cookie_string());
-
-  let response = get_user_info(GetUserInfoRequest {
-    hostname: MidjourneyHostname::Standard,
-    cookie_header: cookie_store.to_cookie_string(),
-  }).await;
-  
-  match response {
-    Err(err) => {
-      error!("Error getting midjourney user info: {:?}", err);
-      // NB: Fall through. Allow us to update the cookies.
-    }
-    Ok(user_info) => {
-      info!("Got midjourney user info: {:?}", user_info);
-      let user_info = MidjourneyUserInfo::from_api_response(user_info);
-      
-      if let Err(err) = mj_creds_manager.replace_user_info(user_info) {
-        error!("Error saving midjourney user info: {:?}", err);
-        // NB: Fall through. Allow us to update the cookies.
-      }
-    }
-  }
-
-  mj_creds_manager.replace_cookie_store(cookie_store)?;
-
-  let result = mj_creds_manager.persist_to_disk();
-
-  if let Err(err) = result {
-    error!("Error persisting midjourney cookies to disk: {:?}", err);
+  let cookies = extract_midjourney_webview_cookies(window)?;
+  // As in ArtCraftX, auth cookies in a cleared login window are the signal;
+  // cookie counts, lengths, and a visit to an OAuth domain are not required.
+  if !cookie_store_has_auth_cookies(&cookies) {
     return Ok(false);
   }
-
-  let event = RefreshAccountStateEvent {
-    provider: Some(GenerationProvider::Midjourney),
-  };
-
-  if let Err(err) = event.send(&app_handle) {
-    error!("Failed to send RefreshAccountStateEvent: {:?}", err); // Fail open
-  }
-
+  credentials.save_login(cookies).await?;
+  RefreshAccountStateEvent { provider: Some(GenerationProvider::Midjourney) }.send_infallible(app);
   Ok(true)
 }
-

crates/desktop/artcraft/src/services/midjourney/windows/open_midjourney_login_window.rs

@@ -1,3 +1,4 @@
+use crate::services::midjourney::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::utils::clear_all_webview_cookies::clear_all_webview_cookies;
 use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
@@ -6,69 +7,37 @@ use anyhow::anyhow;
 use errors::AnyhowResult;
 use once_cell::sync::Lazy;
 use reqwest::Url;
-use std::time::Duration;
 use tauri::{AppHandle, Manager, WebviewUrl, WebviewWindowBuilder};
 
 /// Name of the window
-pub (super) const MIDJOURNEY_LOGIN_WINDOW_NAME: &str = "midjourney_login_window";
+pub(super) const MIDJOURNEY_LOGIN_WINDOW_NAME: &str = "midjourney_login_window";
 
-pub (super) const MIDJOURNEY_START_URL_STR: &str = "https://google.com";
+pub(super) const MIDJOURNEY_START_URL_STR: &str = "https://google.com";
 
-pub (super) static START_URL: Lazy<Url> = Lazy::new(|| {
-  Url::parse(MIDJOURNEY_START_URL_STR).expect("URL should parse")
-});
+pub(super) static START_URL: Lazy<Url> = Lazy::new(|| Url::parse(MIDJOURNEY_START_URL_STR).expect("URL should parse"));
 
-pub (super) const MIDJOURNEY_HOMEPAGE_URL_STR: &str = "https://www.midjourney.com/";
+pub(super) const MIDJOURNEY_HOMEPAGE_URL_STR: &str = "https://www.midjourney.com/";
 
-pub (super) static MIDJOURNEY_HOMEPAGE_URL: Lazy<Url> = Lazy::new(|| {
-  Url::parse(MIDJOURNEY_HOMEPAGE_URL_STR).expect("URL should parse")
-});
+pub(super) static MIDJOURNEY_HOMEPAGE_URL: Lazy<Url> = Lazy::new(|| Url::parse(MIDJOURNEY_HOMEPAGE_URL_STR).expect("URL should parse"));
 
-pub (super) const MIDJOURNEY_LOGIN_URL_STR: &str = "https://www.midjourney.com/auth/signin";
-
-pub (super) static MIDJOURNEY_LOGIN_URL: Lazy<Url> = Lazy::new(|| {
-  Url::parse(MIDJOURNEY_LOGIN_URL_STR).expect("URL should parse")
-});
-
-pub async fn open_midjourney_login_window(
-  app: &AppHandle,
-  app_data_root: &AppDataRoot,
-  mj_creds_manager: &MidjourneyCredentialManager,
-) -> AnyhowResult<()> {
+pub async fn open_midjourney_login_window(app: &AppHandle, app_data_root: &AppDataRoot, mj_creds_manager: &MidjourneyCredentialManager) -> AnyhowResult<()> {
   if app.get_window(MIDJOURNEY_LOGIN_WINDOW_NAME).is_some() {
     return Err(anyhow!("Login window already open"));
   }
 
   let url = WebviewUrl::External(START_URL.clone());
 
-  let window = WebviewWindowBuilder::new(app, MIDJOURNEY_LOGIN_WINDOW_NAME, url)
-      //.user_agent(openai_sora_client::credentials::USER_AGENT)
-      .always_on_top(false)
-      .title("Login to Midjourney")
-      .center()
-      .resizable(true)
-      .visible(true)
-      .closable(true)
-      .min_inner_size(200.0, 800.0)
-      .focused(true)
-      .devtools(true)
-      .build()?;
+  let window = WebviewWindowBuilder::new(app, MIDJOURNEY_LOGIN_WINDOW_NAME, url).user_agent(MIDJOURNEY_USER_AGENT).always_on_top(false).title("Login to Midjourney").center().resizable(true).visible(true).closable(true).min_inner_size(200.0, 800.0).focused(true).devtools(true).build()?;
 
-  let webview = window.get_webview(MIDJOURNEY_LOGIN_WINDOW_NAME)
-      .ok_or_else(|| anyhow!("no webview found"))?;
+  let webview = window.get_webview(MIDJOURNEY_LOGIN_WINDOW_NAME).ok_or_else(|| anyhow!("no webview found"))?;
 
   clear_all_webview_cookies(&webview)?;
 
   webview.navigate(MIDJOURNEY_HOMEPAGE_URL.clone())?;
 
-  // NB: We're starting to get Cloudflare protection screens. Let's try to avoid.
-  tokio::time::sleep(Duration::from_millis(100)).await;
-
-  webview.navigate(MIDJOURNEY_LOGIN_URL.clone())?;
-
   let app_handle = app.clone();
   let app_data_root = app_data_root.clone();
-  let mj_creds_manager= mj_creds_manager.clone();
+  let mj_creds_manager = mj_creds_manager.clone();
 
   let _ = tauri::async_runtime::spawn(async move {
     midjourney_login_window_thread(app_handle, app_data_root, mj_creds_manager).await;

crates/desktop/artcraft/src/services/sora/threads/sora_task_polling/helpers/handle_successful_generations.rs

@@ -1,6 +1,7 @@
 use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
 use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
 use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::utils::auto_download::{auto_download_task_urls, clear_auto_download_checkpoint};
 use crate::core::events::generation_events::generation_failed_event::GenerationFailedEvent;
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
@@ -73,6 +74,10 @@ pub async fn handle_classic_successful_generations(
 
     info!("Task succeeded: {:?}", task_id);
 
+    let local_task = &sqlite_tasks_by_sora_task_id[task_id.as_str()];
+    let urls = generation.items.iter().map(|item| Url::parse(&item.url)).collect::<Result<Vec<_>, _>>()?;
+    auto_download_task_urls(app_handle, local_task, &urls).await?;
+
     let generation_type = match generation.model_type {
       CommonModelType::GptImage1 => GenerationType::Image,
       CommonModelType::Sora2 => GenerationType::Video,
@@ -169,6 +174,7 @@ pub async fn handle_classic_successful_generations(
       }).await?;
 
       if updated {
+        clear_auto_download_checkpoint(app_handle, local_task);
         // If anything breaks with queries, don't spam events.
         let event = GenerationCompleteEvent {
           action: Some(match generation_type {

crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/handle_storyteller_successful_job.rs

@@ -6,6 +6,7 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::task_database::TaskDatabase;
 use crate::core::utils::enum_conversion::generation_provider::to_generation_service_provider;
 use crate::core::utils::enum_conversion::task_type::to_generation_action;
+use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint};
 use super::events::maybe_handle_frontend_caller_notification::maybe_handle_frontend_caller_notification;
 use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
 use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template;
@@ -32,6 +33,13 @@ pub async fn handle_successful_job(
       .as_ref()
       .map(|result| MediaFileToken::new_from_str(&result.entity_token));
 
+  auto_download_task(
+    app_handle, task,
+    job.maybe_result.as_ref().and_then(|result| result.maybe_batch_token.as_ref()),
+    job.maybe_result.as_ref().map(|result| result.media_links.cdn_url.as_str()),
+    None,
+  ).await?;
+
   let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
     db: task_database.get_connection(),
     task_id: &task.id,
@@ -51,6 +59,7 @@ pub async fn handle_successful_job(
     return Ok(()); // If anything breaks with queries, don't spam events.
   }
 
+  clear_auto_download_checkpoint(app_handle, task);
   send_additional_success_events(app_handle, app_env_configs, creds, job, task).await;
 
   let service = to_generation_service_provider(task.provider);

crates/desktop/artcraft/src/services/storyteller/threads/storyteller_task_polling_thread/storyteller_task_polling_thread.rs

@@ -109,7 +109,10 @@ async fn polling_loop(
             TaskStatus::CompleteSuccess => continue, // NB: We're done with this task.
             _ => {}
           }
-          handle_successful_job(app_handle, app_env_configs, creds.as_ref(), job, task, task_database).await?;
+          if let Err(err) = handle_successful_job(app_handle, app_env_configs, creds.as_ref(), job, task, task_database).await {
+            // One failed local download must not delay other completed jobs.
+            error!("Completion for task {} will be retried: {}", task.id.as_str(), err);
+          }
         }
         JobStatusPlus::CompleteFailure => {
           match task.status {

crates/desktop/artcraft/src/services/storyteller/windows/storyteller_billing_window_thread.rs

@@ -21,7 +21,6 @@ use errors::AnyhowResult;
 use log::{error, info};
 use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
 use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
-use midjourney_client::recipes::get_user_info::{get_user_info, GetUserInfoRequest};
 use openai_sora_client::creds::sora_credential_set::SoraCredentialSet;
 use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
 use tauri::{AppHandle, Manager, WebviewWindow};

crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs

@@ -2,6 +2,7 @@ use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCrede
 use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
 use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
 use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
+use crate::core::utils::auto_download::{auto_download_task, clear_auto_download_checkpoint};
 use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
 use crate::core::state::data_dir::app_data_root::AppDataRoot;
 use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
@@ -316,6 +317,8 @@ async fn upload_spz_splat(
     }
   }
 
+  auto_download_task(app_handle, local_task, None, maybe_cdn_url.as_deref(), None).await?;
+
   let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs {
     db: task_database.get_connection(),
     task_id: &local_task.id,
@@ -330,6 +333,7 @@ async fn upload_spz_splat(
     return Ok(()); // If anything breaks with queries, don't spam events.
   }
 
+  clear_auto_download_checkpoint(app_handle, local_task);
   let event = GenerationCompleteEvent {
     //media_file_token: result.media_file_token,
     action: Some(GenerationAction::GenerateGaussian),

crates/lib/browser_emulation/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "browser_emulation"
+edition = "2024"
+version = "0.0.1"
+publish = false
+
+[lib]
+name = "browser_emulation"
+path = "src/lib.rs"
+
+[dependencies]
+# External
+wreq.workspace = true
+wreq-util.workspace = true
+
+[dev-dependencies]
+# None yet

crates/lib/browser_emulation/src/browser_profile.rs

@@ -0,0 +1,225 @@
+use wreq::{Client, ClientBuilder};
+use wreq_util::{Emulation, EmulationOS, EmulationOption};
+
+/// A browser identity to emulate.
+///
+/// The named variants cover the browsers we routinely emulate, defaulting to a
+/// macOS platform (which is what our desktop app and captured cookies use).
+/// [`BrowserProfile::Custom`] is the escape hatch for any other `wreq_util`
+/// emulation, a different OS, or a User-Agent override (useful when a site
+/// advertises a browser version newer than `wreq_util` ships a fingerprint for).
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Default)]
+pub enum BrowserProfile {
+  #[default]
+  Firefox139,
+  Firefox147,
+  Chrome131,
+  Chrome145,
+  Safari18,
+  Custom(CustomBrowserProfile),
+}
+
+/// A fully-specified browser identity for [`BrowserProfile::Custom`].
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct CustomBrowserProfile {
+  /// The `wreq_util` fingerprint to emulate.
+  pub emulation: Emulation,
+
+  /// The operating system to present (drives `sec-ch-ua-platform`, the UA, …).
+  pub os: EmulationOS,
+
+  /// Overrides the User-Agent the emulation would otherwise send. Leave `None`
+  /// to use the emulation's built-in UA.
+  pub maybe_user_agent: Option<String>,
+}
+
+impl BrowserProfile {
+  /// A macOS Safari profile with an explicit User-Agent override.
+  ///
+  /// Use this when the UA must EXACTLY match another client's — e.g. an
+  /// embedded webview that captured the cookies, since Cloudflare's
+  /// `cf_clearance` is bound to the exact User-Agent string.
+  pub fn safari_macos_with_user_agent(user_agent: impl Into<String>) -> Self {
+    Self::Custom(CustomBrowserProfile { emulation: Emulation::Safari18, os: EmulationOS::MacOS, maybe_user_agent: Some(user_agent.into()) })
+  }
+
+  /// The profile that best matches a User-Agent captured from another client
+  /// (typically the login webview that earned the site's cookies), keeping
+  /// that exact UA string.
+  ///
+  /// Use this whenever cookies are replayed: bot-protection clearances
+  /// (Cloudflare `cf_clearance`, DataDome) are bound to the UA and checked
+  /// against the TLS/HTTP fingerprint, so the replaying client should present
+  /// the same UA on the fingerprint of the same browser family. WebKit-based
+  /// UAs (Safari, WKWebView) map to Safari; Chrome/Chromium/Edge to Chrome;
+  /// Firefox to Firefox. Anything else falls back to Safari, since our
+  /// webviews are WebKit.
+  pub fn matching_user_agent(user_agent: impl Into<String>) -> Self {
+    let user_agent = user_agent.into();
+    let emulation = emulation_for_user_agent(&user_agent);
+    Self::Custom(CustomBrowserProfile { emulation, os: os_for_user_agent(&user_agent), maybe_user_agent: Some(user_agent) })
+  }
+
+  /// Build a ready-to-use client carrying this profile's full fingerprint and
+  /// identity headers. This is the usual entry point.
+  pub fn build_client(&self) -> Result<Client, wreq::Error> {
+    self.configure_client_builder(Client::builder()).build()
+  }
+
+  /// Apply this profile to an existing [`ClientBuilder`], so callers can layer
+  /// on their own settings (timeouts, proxies, …) before building.
+  pub fn configure_client_builder(&self, builder: ClientBuilder) -> ClientBuilder {
+    let builder = builder.emulation(self.emulation_option());
+    match self.maybe_user_agent_override() {
+      Some(user_agent) => builder.user_agent(user_agent),
+      None => builder,
+    }
+  }
+
+  /// The `wreq` emulation factory (fingerprint + coherent identity headers)
+  /// for this profile. Pass it to any `wreq` `emulation(...)` method — e.g. to
+  /// emulate at the request level (`RequestBuilder`) or on a websocket upgrade
+  /// (`WebSocketRequestBuilder`) instead of at the client level.
+  ///
+  /// Note: this does NOT include the User-Agent override — that is applied by
+  /// [`Self::configure_client_builder`]. When emulating at the request level,
+  /// apply any override yourself with a `user-agent` header.
+  pub fn emulation_option(&self) -> EmulationOption {
+    EmulationOption::builder().emulation(self.emulation()).emulation_os(self.os()).build()
+  }
+
+  /// The underlying `wreq_util` emulation for this profile.
+  pub fn emulation(&self) -> Emulation {
+    match self {
+      Self::Firefox139 => Emulation::Firefox139,
+      Self::Firefox147 => Emulation::Firefox147,
+      Self::Chrome131 => Emulation::Chrome131,
+      Self::Chrome145 => Emulation::Chrome145,
+      Self::Safari18 => Emulation::Safari18,
+      Self::Custom(custom) => custom.emulation,
+    }
+  }
+
+  /// The operating system this profile presents. Named variants default to
+  /// macOS.
+  pub fn os(&self) -> EmulationOS {
+    match self {
+      Self::Custom(custom) => custom.os,
+      _ => EmulationOS::MacOS,
+    }
+  }
+
+  /// The User-Agent override, if this profile sets one. Named variants use the
+  /// emulation's built-in UA and return `None`.
+  pub fn maybe_user_agent_override(&self) -> Option<&str> {
+    match self {
+      Self::Custom(custom) => custom.maybe_user_agent.as_deref(),
+      _ => None,
+    }
+  }
+
+  /// A short human label for logs and debugging.
+  pub fn label(&self) -> String {
+    match self {
+      Self::Firefox139 => "firefox139".to_string(),
+      Self::Firefox147 => "firefox147".to_string(),
+      Self::Chrome131 => "chrome131".to_string(),
+      Self::Chrome145 => "chrome145".to_string(),
+      Self::Safari18 => "safari18".to_string(),
+      Self::Custom(custom) => {
+        format!("custom({:?} on {:?})", custom.emulation, custom.os)
+      },
+    }
+  }
+}
+
+/// The newest fingerprint we ship for the UA's browser family.
+fn emulation_for_user_agent(user_agent: &str) -> Emulation {
+  let lower = user_agent.to_ascii_lowercase();
+  if lower.contains("firefox/") {
+    Emulation::Firefox147
+  } else if lower.contains("chrome/") || lower.contains("chromium/") || lower.contains("edg/") {
+    Emulation::Chrome145
+  } else {
+    // Safari proper, WKWebView UAs ("AppleWebKit/... Safari/..." without
+    // "Chrome/"), and anything unrecognized: our webviews are WebKit.
+    Emulation::Safari18
+  }
+}
+
+fn os_for_user_agent(user_agent: &str) -> EmulationOS {
+  let lower = user_agent.to_ascii_lowercase();
+  if lower.contains("windows") {
+    EmulationOS::Windows
+  } else if lower.contains("android") {
+    EmulationOS::Android
+  } else if lower.contains("iphone") || lower.contains("ipad") {
+    EmulationOS::IOS
+  } else if lower.contains("linux") || lower.contains("x11") {
+    EmulationOS::Linux
+  } else {
+    EmulationOS::MacOS
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  const WKWEBVIEW_MAC_UA: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko)";
+  const SAFARI_MAC_UA: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.6 Safari/605.1.15";
+  const CHROME_MAC_UA: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/152.0.0.0 Safari/537.36";
+  const FIREFOX_WIN_UA: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:143.0) Gecko/20100101 Firefox/143.0";
+
+  #[test]
+  fn matching_user_agent_picks_family_and_keeps_exact_ua() {
+    let safari = BrowserProfile::matching_user_agent(SAFARI_MAC_UA);
+    assert_eq!(safari.emulation(), Emulation::Safari18);
+    assert_eq!(safari.os(), EmulationOS::MacOS);
+    assert_eq!(safari.maybe_user_agent_override(), Some(SAFARI_MAC_UA));
+
+    // WKWebView says "Safari/..."-less AppleWebKit; still WebKit.
+    assert_eq!(BrowserProfile::matching_user_agent(WKWEBVIEW_MAC_UA).emulation(), Emulation::Safari18);
+
+    // Chrome also says "Safari/537.36" — "Chrome/" must win.
+    let chrome = BrowserProfile::matching_user_agent(CHROME_MAC_UA);
+    assert_eq!(chrome.emulation(), Emulation::Chrome145);
+    assert_eq!(chrome.os(), EmulationOS::MacOS);
+
+    let firefox = BrowserProfile::matching_user_agent(FIREFOX_WIN_UA);
+    assert_eq!(firefox.emulation(), Emulation::Firefox147);
+    assert_eq!(firefox.os(), EmulationOS::Windows);
+  }
+
+  #[test]
+  fn matching_user_agent_builds_a_client() {
+    BrowserProfile::matching_user_agent(CHROME_MAC_UA).build_client().unwrap();
+  }
+
+  #[test]
+  fn default_is_firefox139() {
+    assert_eq!(BrowserProfile::default(), BrowserProfile::Firefox139);
+    assert_eq!(BrowserProfile::default().emulation(), Emulation::Firefox139);
+  }
+
+  #[test]
+  fn named_variants_default_to_macos_and_no_ua_override() {
+    assert_eq!(BrowserProfile::Chrome145.os(), EmulationOS::MacOS);
+    assert_eq!(BrowserProfile::Chrome145.maybe_user_agent_override(), None);
+  }
+
+  #[test]
+  fn custom_carries_os_and_ua_override() {
+    let profile = BrowserProfile::Custom(CustomBrowserProfile { emulation: Emulation::Chrome145, os: EmulationOS::Windows, maybe_user_agent: Some("custom-agent/1.0".to_string()) });
+    assert_eq!(profile.emulation(), Emulation::Chrome145);
+    assert_eq!(profile.os(), EmulationOS::Windows);
+    assert_eq!(profile.maybe_user_agent_override(), Some("custom-agent/1.0"));
+  }
+
+  #[test]
+  fn every_profile_builds_a_client() {
+    for profile in [BrowserProfile::Firefox139, BrowserProfile::Firefox147, BrowserProfile::Chrome131, BrowserProfile::Chrome145, BrowserProfile::Safari18] {
+      profile.build_client().unwrap_or_else(|err| panic!("profile {} failed to build a client: {err}", profile.label()));
+    }
+  }
+}

crates/lib/browser_emulation/src/lib.rs

@@ -0,0 +1,36 @@
+//! browser_emulation
+//!
+//! A small, shared vocabulary of browser identities for anti-bot-sensitive
+//! HTTP clients (Cloudflare-guarded scraping endpoints, websockets, etc.).
+//!
+//! Pick a [`BrowserProfile`] and it configures `wreq` end to end: the TLS/JA3
+//! and HTTP/2 fingerprint, the User-Agent, and the coherent identity headers
+//! (`sec-ch-ua`, `accept`, `accept-encoding`, …) that a real browser sends.
+//! These MUST live on the client/connection, not on individual requests, so
+//! the primary entry point builds a configured client:
+//!
+//! ```ignore
+//! let client = BrowserProfile::Firefox139.build_client()?;
+//! // Every request from `client` now carries Firefox 139's fingerprint.
+//! ```
+//!
+//! When the cookies you send were captured from a specific browser (e.g. a
+//! Cloudflare `cf_clearance` cookie, which is bound to the capturing browser's
+//! User-Agent), pick the profile that matches that browser so the cookie
+//! validates.
+
+// Never allow these
+#![forbid(private_bounds)]
+#![forbid(private_interfaces)]
+#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things
+
+// Okay to toggle
+#![forbid(unreachable_patterns)]
+#![forbid(unused_imports)]
+#![forbid(unused_mut)]
+#![forbid(unused_variables)]
+// Always allow
+#![allow(dead_code)]
+#![allow(non_snake_case)]
+
+pub mod browser_profile;

crates/lib/cookie_store_wrapper/Cargo.toml

@@ -0,0 +1,37 @@
+[package]
+name = "cookie_store_wrapper"
+edition = "2024"
+version = "0.0.2"
+publish = false
+
+[lib]
+name = "cookie_store_wrapper"
+path = "src/lib.rs"
+
+[features]
+default = []
+# Implements the wreq `CookieStore` trait on `SharedCookieStore` so a wreq
+# client can read from and write Set-Cookie responses into a shared store.
+wreq = ["dep:wreq"]
+
+[dependencies]
+# External
+chrono.workspace = true
+log.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+url.workspace = true
+
+# NB: For converting cookie expiry timestamps; the `cookie` crate (via
+# `cookie_store`) already depends on `time`, so this adds no new weight.
+time = "0.3"
+wreq = { workspace = true, optional = true }
+
+# NB: The RFC 6265 cookie store implementation this crate wraps. `serde_json`
+# enables cookie serialization; we skip `public_suffix` to keep dependencies
+# minimal.
+cookie_store = { version = "0.21.1", default-features = false, features = ["serde_json"] }
+
+[dev-dependencies]
+serde_derive.workspace = true
+toml.workspace = true

crates/lib/cookie_store_wrapper/src/change_log.rs

@@ -0,0 +1,80 @@
+use std::collections::VecDeque;
+use std::time::SystemTime;
+
+/// Oldest records are dropped once the log exceeds this many entries.
+pub const MAX_CHANGE_LOG_RECORDS: usize = 256;
+
+/// A bounded, in-memory history of cookie store mutations, for debugging
+/// auth flows ("did the provider rotate the session cookie?", "when did we
+/// lose it?"). Values are intentionally not recorded — only names, domains,
+/// and what happened.
+#[derive(Clone, Debug, Default)]
+pub struct CookieChangeLog {
+  records: VecDeque<CookieChangeRecord>,
+  next_sequence_number: u64,
+}
+
+#[derive(Clone, Debug)]
+pub struct CookieChangeRecord {
+  pub sequence_number: u64,
+  pub timestamp: SystemTime,
+  pub action: CookieChangeAction,
+  pub cookie_name: String,
+  pub maybe_domain: Option<String>,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum CookieChangeAction {
+  /// A new cookie was stored.
+  Inserted,
+  /// An existing cookie's value was replaced.
+  Updated,
+  /// A server sent an expired Set-Cookie, deleting the stored cookie.
+  ExpiredByServer,
+  /// A cookie was refused (domain mismatch, malformed, etc.).
+  Rejected,
+  /// The whole store was cleared.
+  Cleared,
+}
+
+impl CookieChangeLog {
+  pub fn record(&mut self, action: CookieChangeAction, cookie_name: &str, maybe_domain: Option<String>) {
+    self.records.push_back(CookieChangeRecord { sequence_number: self.next_sequence_number, timestamp: SystemTime::now(), action, cookie_name: cookie_name.to_owned(), maybe_domain });
+    self.next_sequence_number += 1;
+    while self.records.len() > MAX_CHANGE_LOG_RECORDS {
+      self.records.pop_front();
+    }
+  }
+
+  /// Records in chronological order, oldest first.
+  pub fn iter(&self) -> impl Iterator<Item = &CookieChangeRecord> {
+    self.records.iter()
+  }
+
+  pub fn len(&self) -> usize {
+    self.records.len()
+  }
+
+  pub fn is_empty(&self) -> bool {
+    self.records.is_empty()
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn log_is_bounded_and_keeps_newest_records() {
+    let mut log = CookieChangeLog::default();
+    for index in 0..(MAX_CHANGE_LOG_RECORDS + 10) {
+      log.record(CookieChangeAction::Inserted, &format!("cookie_{index}"), None);
+    }
+
+    assert_eq!(log.len(), MAX_CHANGE_LOG_RECORDS);
+    let first = log.iter().next().unwrap();
+    assert_eq!(first.sequence_number, 10);
+    let last = log.iter().last().unwrap();
+    assert_eq!(last.cookie_name, format!("cookie_{}", MAX_CHANGE_LOG_RECORDS + 9));
+  }
+}

crates/lib/cookie_store_wrapper/src/cookie_store.rs

@@ -0,0 +1,329 @@
+use crate::change_log::{CookieChangeAction, CookieChangeLog};
+use log::warn;
+use ::cookie_store::{CookieError, CookieStore as CookieStoreRfc, RawCookie, StoreAction};
+use url::Url;
+
+/// A cookie jar with RFC 6265 storage semantics.
+///
+/// Cookies are keyed by `(domain, path, name)` — not just name — and a repeat
+/// `Set-Cookie` for the same key replaces the old value, while an expired
+/// `Set-Cookie` deletes it (the standard server-side removal mechanism).
+/// Every mutation is recorded in a bounded [`CookieChangeLog`] for debugging.
+///
+/// This is a plain value type; wrap it in
+/// [`crate::shared_cookie_store::SharedCookieStore`] to share it between
+/// threads.
+#[derive(Clone, Debug, Default)]
+pub struct CookieStore {
+  store: CookieStoreRfc,
+  change_log: CookieChangeLog,
+}
+
+/// A cookie captured from a browser or webview, with as much metadata as the
+/// browser reported.
+#[derive(Clone, Debug)]
+pub struct CapturedCookie {
+  pub name: String,
+  pub value: String,
+
+  /// Domain as reported by the browser. May carry a leading dot
+  /// (e.g. `.storyteller.ai`), which marks a parent-domain cookie.
+  pub maybe_domain: Option<String>,
+
+  pub maybe_path: Option<String>,
+}
+
+impl CookieStore {
+  pub fn empty() -> Self {
+    Self::default()
+  }
+
+  /// Parse a raw `Cookie:` request-header string (`"a=1; b=2"`) into a new
+  /// store of host-only cookies for `source_url`'s host. This is the entry
+  /// point for hand-entered cookie headers, where per-cookie attributes are
+  /// not known.
+  pub fn from_cookie_header(cookie_header: &str, source_url: &Url) -> Self {
+    let mut store = Self::empty();
+    store.insert_cookie_header(cookie_header, source_url);
+    store
+  }
+
+  pub(crate) fn from_rfc_store(store: CookieStoreRfc) -> Self {
+    Self { store, change_log: CookieChangeLog::default() }
+  }
+
+  /// Insert a bare name/value cookie as a host-only cookie for `source_url`'s
+  /// host. Returns whether the cookie was stored (or expired an existing one).
+  pub fn insert_named(&mut self, name: &str, value: &str, source_url: &Url) -> bool {
+    let raw = RawCookie::new(name.to_owned(), value.to_owned());
+    self.insert_raw(&raw, source_url)
+  }
+
+  /// Insert a cookie captured from a browser or webview.
+  ///
+  /// When the capture carries a domain, the cookie is stored as a
+  /// domain-suffix cookie for it (so it matches subdomains, mirroring how the
+  /// browser scoped it). Without a domain it falls back to a host-only cookie
+  /// for `fallback_source_url`.
+  pub fn insert_captured(&mut self, captured: CapturedCookie, fallback_source_url: &Url) -> bool {
+    let maybe_domain = captured.maybe_domain.as_deref().map(|domain| domain.trim_start_matches('.').to_ascii_lowercase()).filter(|domain| !domain.is_empty());
+
+    let Some(domain) = maybe_domain else {
+      return self.insert_named(&captured.name, &captured.value, fallback_source_url);
+    };
+
+    let source_url = match Url::parse(&format!("https://{domain}/")) {
+      Ok(url) => url,
+      Err(err) => {
+        warn!("Cannot build source URL for captured cookie {:?} (domain {:?}): {:?}", captured.name, domain, err);
+        return false;
+      },
+    };
+
+    let mut raw = RawCookie::new(captured.name, captured.value);
+    raw.set_domain(domain);
+    raw.set_path(captured.maybe_path.unwrap_or_else(|| "/".to_owned()));
+    self.insert_raw(&raw, &source_url)
+  }
+
+  /// Insert every `name=value` pair of a raw `Cookie:` request-header string
+  /// as a host-only cookie for `source_url`'s host. Returns how many cookies
+  /// were stored.
+  pub fn insert_cookie_header(&mut self, cookie_header: &str, source_url: &Url) -> usize {
+    cookie_header.split(';').filter_map(|pair| pair.trim().split_once('=')).filter(|(name, _)| !name.trim().is_empty()).map(|(name, value)| (name.trim().to_owned(), value.trim().to_owned())).collect::<Vec<(String, String)>>().into_iter().filter(|(name, value)| self.insert_named(name, value, source_url)).count()
+  }
+
+  /// Apply one `Set-Cookie` header value received from `request_url`,
+  /// following the RFC 6265 storage model (replace / expire / reject).
+  pub fn apply_set_cookie_header(&mut self, set_cookie_value: &str, request_url: &Url) -> bool {
+    let cookie_name = set_cookie_value.split('=').next().unwrap_or(set_cookie_value).trim().to_owned();
+    let result = self.store.parse(set_cookie_value, request_url);
+    self.record_insert_result(result, &cookie_name, request_url)
+  }
+
+  /// Apply several `Set-Cookie` header values received from `request_url`.
+  /// Returns how many were accepted.
+  pub fn apply_set_cookie_headers<'a>(&mut self, set_cookie_values: impl Iterator<Item = &'a str>, request_url: &Url) -> usize {
+    set_cookie_values.filter(|value| self.apply_set_cookie_header(value, request_url)).count()
+  }
+
+  pub fn clear_all(&mut self) {
+    self.store.clear();
+    self.change_log.record(CookieChangeAction::Cleared, "*", None);
+  }
+
+  /// Whether an unexpired cookie with this name exists on any domain.
+  pub fn has_cookie(&self, name: &str) -> bool {
+    self.store.iter_unexpired().any(|cookie| cookie.name() == name)
+  }
+
+  /// The value of the first unexpired cookie with this name, on any domain.
+  pub fn get_cookie_value(&self, name: &str) -> Option<&str> {
+    self.store.iter_unexpired().find(|cookie| cookie.name() == name).map(|cookie| cookie.value())
+  }
+
+  /// `(name, value)` pairs of every unexpired cookie, across all domains.
+  pub fn iter_name_values(&self) -> impl Iterator<Item = (&str, &str)> {
+    self.store.iter_unexpired().map(|cookie| (cookie.name(), cookie.value()))
+  }
+
+  /// Names of every unexpired cookie, across all domains.
+  pub fn cookie_names(&self) -> Vec<String> {
+    self.iter_name_values().map(|(name, _)| name.to_owned()).collect()
+  }
+
+  pub fn len(&self) -> usize {
+    self.store.iter_unexpired().count()
+  }
+
+  pub fn is_empty(&self) -> bool {
+    self.store.iter_unexpired().next().is_none()
+  }
+
+  /// NB: Just use this as a heuristic, and do not call it in a loop.
+  pub fn calculate_approx_cookie_character_length(&self) -> usize {
+    const COOKIE_SEP_LENGTH: usize = 2; // '=' and ';'
+    self.iter_name_values().map(|(name, value)| name.len() + value.len() + COOKIE_SEP_LENGTH).sum()
+  }
+
+  /// Every unexpired cookie joined into one `Cookie:` header string,
+  /// regardless of domain or path.
+  ///
+  /// This is a blunt tool for scraping clients that captured cookies for one
+  /// known site and want to send them all back. It performs no request
+  /// matching; two same-named cookies on different domains both appear.
+  /// Prefer [`Self::cookie_header_for_url`] when the request URL is known.
+  pub fn to_cookie_string(&self) -> String {
+    self.iter_name_values().map(|(name, value)| format!("{}={}", name, value)).collect::<Vec<String>>().join("; ")
+  }
+
+  /// A `Cookie:` header value containing only the cookies that RFC 6265
+  /// domain/path/secure matching would send to `request_url`, or `None` when
+  /// no cookies match.
+  pub fn cookie_header_for_url(&self, request_url: &Url) -> Option<String> {
+    let header = self.store.get_request_values(request_url).map(|(name, value)| format!("{}={}", name, value)).collect::<Vec<String>>().join("; ");
+    if header.is_empty() { None } else { Some(header) }
+  }
+
+  /// The recent mutation history, for debugging.
+  pub fn change_log(&self) -> &CookieChangeLog {
+    &self.change_log
+  }
+
+  pub(crate) fn rfc_store(&self) -> &CookieStoreRfc {
+    &self.store
+  }
+
+  pub(crate) fn insert_raw(&mut self, raw: &RawCookie<'static>, source_url: &Url) -> bool {
+    let cookie_name = raw.name().to_owned();
+    let result = self.store.insert_raw(raw, source_url);
+    self.record_insert_result(result, &cookie_name, source_url)
+  }
+
+  fn record_insert_result(&mut self, result: Result<StoreAction, CookieError>, cookie_name: &str, request_url: &Url) -> bool {
+    let maybe_domain = request_url.host_str().map(str::to_owned);
+    match result {
+      Ok(StoreAction::Inserted) => {
+        self.change_log.record(CookieChangeAction::Inserted, cookie_name, maybe_domain);
+        true
+      },
+      Ok(StoreAction::UpdatedExisting) => {
+        self.change_log.record(CookieChangeAction::Updated, cookie_name, maybe_domain);
+        true
+      },
+      Ok(StoreAction::ExpiredExisting) => {
+        self.change_log.record(CookieChangeAction::ExpiredByServer, cookie_name, maybe_domain);
+        true
+      },
+      Err(err) => {
+        warn!("Rejected cookie {:?} from {}: {:?}", cookie_name, request_url, err);
+        self.change_log.record(CookieChangeAction::Rejected, cookie_name, maybe_domain);
+        false
+      },
+    }
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use crate::change_log::CookieChangeAction;
+
+  const SITE_URL: &str = "https://www.example.com/";
+
+  mod insert_tests {
+    use super::*;
+
+    #[test]
+    fn insert_named_stores_host_only_cookie() {
+      let mut store = CookieStore::empty();
+      assert!(store.insert_named("session", "abc123", &site_url()));
+
+      assert!(store.has_cookie("session"));
+      assert_eq!(store.get_cookie_value("session"), Some("abc123"));
+      assert_eq!(store.len(), 1);
+    }
+
+    #[test]
+    fn insert_named_replaces_existing_value() {
+      let mut store = CookieStore::empty();
+      store.insert_named("session", "old", &site_url());
+      store.insert_named("session", "new", &site_url());
+
+      assert_eq!(store.len(), 1);
+      assert_eq!(store.get_cookie_value("session"), Some("new"));
+    }
+
+    #[test]
+    fn insert_captured_with_parent_domain_matches_subdomains() {
+      let mut store = CookieStore::empty();
+      let captured = CapturedCookie { name: "auth".to_string(), value: "token".to_string(), maybe_domain: Some(".example.com".to_string()), maybe_path: None };
+      assert!(store.insert_captured(captured, &site_url()));
+
+      let subdomain_url = Url::parse("https://api.example.com/v1/things").unwrap();
+      assert_eq!(store.cookie_header_for_url(&subdomain_url), Some("auth=token".to_string()));
+    }
+
+    #[test]
+    fn insert_captured_without_domain_falls_back_to_source_url() {
+      let mut store = CookieStore::empty();
+      let captured = CapturedCookie { name: "sid".to_string(), value: "xyz".to_string(), maybe_domain: None, maybe_path: None };
+      assert!(store.insert_captured(captured, &site_url()));
+
+      assert_eq!(store.cookie_header_for_url(&site_url()), Some("sid=xyz".to_string()));
+    }
+  }
+
+  mod set_cookie_header_tests {
+    use super::*;
+
+    #[test]
+    fn apply_set_cookie_header_stores_cookie() {
+      let mut store = CookieStore::empty();
+      assert!(store.apply_set_cookie_header("session=abc; Path=/; HttpOnly", &site_url()));
+      assert_eq!(store.get_cookie_value("session"), Some("abc"));
+    }
+
+    #[test]
+    fn expired_set_cookie_deletes_existing_cookie() {
+      let mut store = CookieStore::empty();
+      store.apply_set_cookie_header("session=abc; Path=/", &site_url());
+      assert!(store.has_cookie("session"));
+
+      let accepted = store.apply_set_cookie_header("session=gone; Path=/; Expires=Thu, 01 Jan 1970 00:00:00 GMT", &site_url());
+      assert!(accepted);
+      assert!(!store.has_cookie("session"));
+    }
+
+    #[test]
+    fn set_cookie_for_mismatched_domain_is_rejected() {
+      let mut store = CookieStore::empty();
+      let accepted = store.apply_set_cookie_header("evil=1; Domain=other.com", &site_url());
+      assert!(!accepted);
+      assert!(!store.has_cookie("evil"));
+    }
+  }
+
+  mod header_rendering_tests {
+    use super::*;
+
+    #[test]
+    fn to_cookie_string_joins_all_cookies() {
+      let mut store = CookieStore::empty();
+      store.insert_named("a", "1", &site_url());
+      store.insert_named("b", "2", &site_url());
+
+      let header = store.to_cookie_string();
+      let mut parts = header.split("; ").collect::<Vec<_>>();
+      parts.sort();
+      assert_eq!(parts, vec!["a=1", "b=2"]);
+    }
+
+    #[test]
+    fn cookie_header_for_url_excludes_other_hosts() {
+      let mut store = CookieStore::empty();
+      store.insert_named("mine", "1", &site_url());
+      store.insert_named("theirs", "2", &Url::parse("https://other.example.org/").unwrap());
+
+      assert_eq!(store.cookie_header_for_url(&site_url()), Some("mine=1".to_string()));
+    }
+  }
+
+  mod change_log_tests {
+    use super::*;
+
+    #[test]
+    fn mutations_are_recorded() {
+      let mut store = CookieStore::empty();
+      store.insert_named("session", "a", &site_url());
+      store.insert_named("session", "b", &site_url());
+
+      let actions = store.change_log().iter().map(|record| record.action).collect::<Vec<_>>();
+      assert_eq!(actions, vec![CookieChangeAction::Inserted, CookieChangeAction::Updated]);
+    }
+  }
+
+  fn site_url() -> Url {
+    Url::parse(SITE_URL).unwrap()
+  }
+}

crates/lib/cookie_store_wrapper/src/lib.rs

@@ -0,0 +1,36 @@
+//! cookie_store_wrapper
+//!
+//! Cookie jar for scraping-style API clients.
+//!
+//! Wraps the RFC 6265 `cookie_store` crate so stored cookies follow real HTTP
+//! semantics: domain and path matching, replacement on repeat Set-Cookie, and
+//! server-driven expiry. On top of that this crate adds Serde support (embed a
+//! [`cookie_store::CookieStore`] directly in your own serializable structs), a
+//! bounded change log for debugging auth flows, and a cheaply clonable
+//! thread-safe handle ([`shared_cookie_store::SharedCookieStore`]).
+//!
+//! With the `wreq` feature enabled, `SharedCookieStore` implements wreq's
+//! `CookieStore` trait, so a wreq client built with it as the cookie provider
+//! transparently applies every `Set-Cookie` response to the shared store.
+
+// Never allow these
+#![forbid(private_bounds)]
+#![forbid(private_interfaces)]
+#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things
+
+// Okay to toggle
+#![forbid(unreachable_patterns)]
+#![forbid(unused_imports)]
+#![forbid(unused_mut)]
+#![forbid(unused_variables)]
+// Always allow
+#![allow(dead_code)]
+#![allow(non_snake_case)]
+
+pub mod change_log;
+pub mod cookie_store;
+pub mod serde_support;
+pub mod shared_cookie_store;
+
+#[cfg(feature = "wreq")]
+pub mod wreq_cookie_provider;

crates/lib/cookie_store_wrapper/src/serde_support.rs

@@ -0,0 +1,312 @@
+//! Serde support for [`CookieStore`] and [`SharedCookieStore`].
+//!
+//! Both types implement `Serialize` and `Deserialize` directly, so callers
+//! can embed a cookie store in their own serializable structs (JSON state
+//! files, TOML credential files, etc.) without a conversion type:
+//!
+//! ```ignore
+//! #[derive(Serialize, Deserialize)]
+//! struct CookieCredential {
+//!   updated_at: Option<DateTime<Utc>>,
+//!   cookies: CookieStore,  // NB: keep array-of-table fields LAST for TOML
+//! }
+//! ```
+//!
+//! The wire shape is a flat list of cookies, chosen to stay human-editable
+//! when it lands in a TOML file someone opens in vim:
+//!
+//! ```toml
+//! [[cookie.cookies]]
+//! name = "session"
+//! value = "abc123"
+//! domain = "example.com"
+//! secure = true
+//! http_only = true
+//! expires = "2027-01-01T00:00:00+00:00"
+//! ```
+//!
+//! Defaulted attributes (`path = "/"`, false flags, session-cookie expiry)
+//! are omitted on write but accepted on read. Expired cookies are dropped on
+//! both save and load. The change log is debug state and is not serialized.
+
+use crate::cookie_store::CookieStore;
+use crate::shared_cookie_store::SharedCookieStore;
+use ::cookie_store::Cookie as CookieRfc;
+use ::cookie_store::{CookieDomain, CookieExpiration, RawCookie};
+use chrono::{DateTime, SecondsFormat, Utc};
+use log::warn;
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+use std::time::SystemTime;
+use url::Url;
+
+const DEFAULT_PATH: &str = "/";
+
+/// The wire shape of one cookie. Public so callers can build tooling around
+/// the stored format, but normal use is via `CookieStore`'s serde impls.
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
+pub struct SerializedCookie {
+  pub name: String,
+  pub value: String,
+
+  /// The domain the cookie belongs to, without a leading dot. Subdomains
+  /// also match unless `host_only` is set.
+  pub domain: String,
+
+  /// Restrict the cookie to exactly `domain` (no subdomains).
+  #[serde(default, skip_serializing_if = "is_false")]
+  pub host_only: bool,
+
+  #[serde(default = "default_path", skip_serializing_if = "is_default_path")]
+  pub path: String,
+
+  #[serde(default, skip_serializing_if = "is_false")]
+  pub secure: bool,
+
+  #[serde(default, skip_serializing_if = "is_false")]
+  pub http_only: bool,
+
+  /// RFC 3339 expiry time. Omitted for session cookies (which never expire
+  /// on disk).
+  #[serde(default, skip_serializing_if = "Option::is_none")]
+  pub expires: Option<String>,
+}
+
+impl Serialize for CookieStore {
+  fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+    serializer.collect_seq(self.rfc_store().iter_unexpired().filter_map(serialized_cookie_from_rfc))
+  }
+}
+
+impl<'de> Deserialize<'de> for CookieStore {
+  fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
+    let cookies = Vec::<SerializedCookie>::deserialize(deserializer)?;
+    let mut store = CookieStore::empty();
+    for cookie in cookies {
+      insert_serialized_cookie(&mut store, cookie);
+    }
+    Ok(store)
+  }
+}
+
+/// Serializes as a point-in-time snapshot of the shared store.
+impl Serialize for SharedCookieStore {
+  fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+    self.with_read(|store| store.serialize(serializer))
+  }
+}
+
+/// Deserializes into a fresh handle, unrelated to any existing clones.
+impl<'de> Deserialize<'de> for SharedCookieStore {
+  fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
+    Ok(SharedCookieStore::new(CookieStore::deserialize(deserializer)?))
+  }
+}
+
+/// Compares the cookies (name, value, and attributes), ignoring the change
+/// log and insertion order.
+impl PartialEq for CookieStore {
+  fn eq(&self, other: &Self) -> bool {
+    sorted_serialized_cookies(self) == sorted_serialized_cookies(other)
+  }
+}
+
+impl Eq for CookieStore {}
+
+fn serialized_cookie_from_rfc(cookie: &CookieRfc<'static>) -> Option<SerializedCookie> {
+  let (domain, host_only) = match &cookie.domain {
+    CookieDomain::HostOnly(host) => (host.clone(), true),
+    CookieDomain::Suffix(suffix) => (suffix.clone(), false),
+    other => {
+      warn!("Cannot serialize cookie {:?} with domain {:?}; skipping.", cookie.name(), other);
+      return None;
+    },
+  };
+
+  let expires = match &cookie.expires {
+    CookieExpiration::AtUtc(at_utc) => {
+      let system_time = SystemTime::from(*at_utc);
+      Some(DateTime::<Utc>::from(system_time).to_rfc3339_opts(SecondsFormat::Secs, false))
+    },
+    CookieExpiration::SessionEnd => None,
+  };
+
+  Some(SerializedCookie { name: cookie.name().to_owned(), value: cookie.value().to_owned(), domain, host_only, path: (*cookie.path).to_owned(), secure: cookie.secure().unwrap_or(false), http_only: cookie.http_only().unwrap_or(false), expires })
+}
+
+fn insert_serialized_cookie(store: &mut CookieStore, cookie: SerializedCookie) {
+  let source_url = match Url::parse(&format!("https://{}/", cookie.domain)) {
+    Ok(url) => url,
+    Err(err) => {
+      warn!("Cannot restore cookie {:?} with domain {:?}: {:?}", cookie.name, cookie.domain, err);
+      return;
+    },
+  };
+
+  let mut raw = RawCookie::new(cookie.name, cookie.value);
+  raw.set_path(cookie.path);
+  raw.set_secure(cookie.secure);
+  raw.set_http_only(cookie.http_only);
+  if !cookie.host_only {
+    raw.set_domain(cookie.domain);
+  }
+  if let Some(expires) = &cookie.expires {
+    match DateTime::parse_from_rfc3339(expires) {
+      Ok(at) => {
+        let system_time = SystemTime::from(at.with_timezone(&Utc));
+        raw.set_expires(time::OffsetDateTime::from(system_time));
+      },
+      Err(err) => {
+        warn!("Cookie {:?} has unparseable expires {:?} ({:?}); treating as session cookie.", raw.name(), expires, err);
+      },
+    }
+  }
+
+  // NB: An expired cookie is rejected here (with a warning), which is the
+  // intended load behavior.
+  store.insert_raw(&raw, &source_url);
+}
+
+fn sorted_serialized_cookies(store: &CookieStore) -> Vec<SerializedCookie> {
+  let mut cookies = store.rfc_store().iter_unexpired().filter_map(serialized_cookie_from_rfc).collect::<Vec<SerializedCookie>>();
+  cookies.sort_by(|a, b| (&a.domain, &a.path, &a.name).cmp(&(&b.domain, &b.path, &b.name)));
+  cookies
+}
+
+fn default_path() -> String {
+  DEFAULT_PATH.to_owned()
+}
+
+fn is_default_path(path: &str) -> bool {
+  path == DEFAULT_PATH
+}
+
+fn is_false(value: &bool) -> bool {
+  !value
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use serde_derive::{Deserialize, Serialize};
+
+  #[test]
+  fn cookie_store_round_trips_through_json() {
+    let store = example_store();
+
+    let serialized = serde_json::to_string(&store).unwrap();
+    let restored: CookieStore = serde_json::from_str(&serialized).unwrap();
+
+    assert_eq!(restored.len(), 3);
+    assert_eq!(restored.get_cookie_value("session"), Some("abc123"));
+    assert_eq!(restored.get_cookie_value("pref"), Some("dark"));
+    assert_eq!(restored.get_cookie_value("keeper"), Some("forever"));
+    assert_eq!(restored, store);
+  }
+
+  #[test]
+  fn cookie_store_nests_inside_caller_toml_structs() {
+    #[derive(Serialize, Deserialize)]
+    struct CredentialFile {
+      service: String,
+      // NB: Array-of-table fields go last so TOML stays valid.
+      cookies: CookieStore,
+    }
+
+    let file = CredentialFile { service: "midjourney_cookies".to_string(), cookies: example_store() };
+
+    let serialized = toml::to_string(&file).unwrap();
+    let restored: CredentialFile = toml::from_str(&serialized).unwrap();
+
+    assert_eq!(restored.service, "midjourney_cookies");
+    assert_eq!(restored.cookies, file.cookies);
+  }
+
+  #[test]
+  fn toml_output_is_human_editable() {
+    #[derive(Serialize)]
+    struct CredentialFile {
+      cookies: CookieStore,
+    }
+
+    let site_url = Url::parse("https://www.example.com/").unwrap();
+    let mut store = CookieStore::empty();
+    store.apply_set_cookie_header("session=abc123; Secure; HttpOnly; Path=/", &site_url);
+
+    let serialized = toml::to_string(&CredentialFile { cookies: store }).unwrap();
+
+    // The whole cookie reads as a flat, obvious table.
+    let expected = "\
+[[cookies]]
+name = \"session\"
+value = \"abc123\"
+domain = \"www.example.com\"
+host_only = true
+secure = true
+http_only = true
+";
+    assert_eq!(serialized, expected);
+  }
+
+  #[test]
+  fn hand_written_toml_with_defaults_parses() {
+    #[derive(Deserialize)]
+    struct CredentialFile {
+      cookies: CookieStore,
+    }
+
+    let hand_written = r#"
+      [[cookies]]
+      name = "session"
+      value = "abc123"
+      domain = "example.com"
+    "#;
+
+    let file: CredentialFile = toml::from_str(hand_written).unwrap();
+    assert_eq!(file.cookies.get_cookie_value("session"), Some("abc123"));
+
+    // Domain cookies (the default) match subdomains.
+    let subdomain_url = Url::parse("https://api.example.com/").unwrap();
+    assert_eq!(file.cookies.cookie_header_for_url(&subdomain_url), Some("session=abc123".to_string()),);
+  }
+
+  #[test]
+  fn expires_round_trips_and_expired_cookies_are_dropped() {
+    let json = r#"[
+      {"name": "fresh", "value": "1", "domain": "example.com", "expires": "2100-01-01T00:00:00+00:00"},
+      {"name": "stale", "value": "2", "domain": "example.com", "expires": "2001-01-01T00:00:00+00:00"}
+    ]"#;
+
+    let restored: CookieStore = serde_json::from_str(json).unwrap();
+    assert!(restored.has_cookie("fresh"));
+    assert!(!restored.has_cookie("stale"));
+
+    let reserialized = serde_json::to_string(&restored).unwrap();
+    assert!(reserialized.contains("\"expires\":\"2100-01-01T00:00:00+00:00\""));
+  }
+
+  #[test]
+  fn shared_cookie_store_round_trips_through_json() {
+    let shared = SharedCookieStore::new(example_store());
+
+    let serialized = serde_json::to_string(&shared).unwrap();
+    let restored: SharedCookieStore = serde_json::from_str(&serialized).unwrap();
+
+    assert!(restored.has_cookie("session"));
+  }
+
+  #[test]
+  fn legacy_version_1_data_fails_to_parse() {
+    let legacy_json = r#"{"cookies":[{"name":"session","value":"abc"}]}"#;
+    let result = serde_json::from_str::<CookieStore>(legacy_json);
+    assert!(result.is_err());
+  }
+
+  fn example_store() -> CookieStore {
+    let site_url = Url::parse("https://www.example.com/").unwrap();
+    let mut store = CookieStore::empty();
+    store.insert_named("session", "abc123", &site_url);
+    store.apply_set_cookie_header("pref=dark; Domain=example.com; Path=/", &site_url);
+    store.apply_set_cookie_header("keeper=forever; Path=/; Expires=Tue, 03 Aug 2100 00:38:37 GMT", &site_url);
+    store
+  }
+}