artcraft-native: generate with Midjourney or fal (provider choice) (#2115)

Images can run on Midjourney (own account, via a sign-in window) or fal (own API key), chosen from the model selector's Select Provider submenu. Jobs finish into the library. The Midjourney credential manager and the router prompt helper move into shared crates used by both desktop apps.
This commit is contained in:
BFlat authored and GitHub committed 2026-10-11 20:47:59 +07:00
1 parent 28a39a0de4
commit d4e4ac69bb
39 files changed
+2169 -464

No files matched your search

Generated
+12
View File
@@ -487,8 +487,10 @@ dependencies = [
"arboard",
"artcraft_api_defs",
"artcraft_client",
"artcraft_router",
"chrono",
"cookie",
"cookie_store_wrapper",
"directories",
"eframe",
"egui",
@@ -496,8 +498,10 @@ dependencies = [
"egui_kittest",
"enums",
"env_logger",
"fal_client",
"image",
"log",
"midjourney_client",
"open",
"reqwest 0.12.23",
"resvg",
@@ -505,9 +509,12 @@ dependencies = [
"rodio",
"serde",
"serde_json",
"tempfile",
"tokens",
"tokio",
"url",
"uuid 1.22.0",
"wry",
]
[[package]]
@@ -5009,7 +5016,9 @@ name = "midjourney_client"
version = "0.0.1"
dependencies = [
"anyhow",
"base64 0.22.1",
"browser_emulation",
"chrono",
"ciborium",
"cloudflare_errors",
"cookie_store_wrapper",
@@ -5020,10 +5029,13 @@ dependencies = [
"log",
"serde",
"serde_json",
"tempfile",
"test_data",
"test_utils",
"tokio",
"toml 1.1.8+spec-1.1.0",
"url",
"uuid_utils",
"wreq",
"wreq-util",
]
+9
View File
@@ -13,6 +13,10 @@ path = "src/main.rs"
# Internal (Tauri-free) API crates shared with the Tauri app.
artcraft_api_defs.workspace = true
artcraft_client.workspace = true
artcraft_router.workspace = true
cookie_store_wrapper.workspace = true
fal_client.workspace = true
midjourney_client.workspace = true
tokens.workspace = true
# External
@@ -34,11 +38,16 @@ rfd = "0.16"
rodio = { version = "0.20", default-features = false, features = ["symphonia-all"] }
serde.workspace = true
serde_json.workspace = true
tempfile.workspace = true
tokio = { workspace = true, features = ["rt-multi-thread", "sync", "time", "fs"] }
url.workspace = true
uuid = { workspace = true, features = ["v4"] }
# Pasting images from the clipboard (egui's own clipboard support is text only).
arboard = { version = "3", default-features = false, features = ["image-data"] }
# The Midjourney sign-in window (its login cookies can only be captured inside a browser).
wry = "0.55"
[dev-dependencies]
egui_kittest = "0.36"
+13 -5
View File
@@ -1,7 +1,8 @@
# ArtCraft Native
ArtCraft's desktop app in pure Rust: an [egui](https://github.com/emilk/egui) UI over the same
ArtCraft API client crates the Tauri app uses. There is no Tauri, no webview and no JavaScript.
ArtCraft API client crates the Tauri app uses. There is no Tauri and no JavaScript; the only
webview is the Midjourney sign-in window, which exists to capture Midjourney's login cookies.
The look follows the brutalist webapp (`artcraft-services/frontend/apps/artcraft-webapp`). That
means near-black surfaces with white/15 hairlines and 3 px corners on controls. Type is Archivo for
@@ -33,11 +34,12 @@ The workspace's default member is the Tauri app, so pass `-p artcraft_native`.
| `backend::session` | Credentials: own file, Tauri cookie import, env config |
| `backend::media_cache` | Thumbnails and previews: downloaded and decoded off-thread, evicted when unused |
| `backend::audio` | One clip at a time: reference previews, feed tracks (pause, seek, volume), waveforms |
| `backend::providers` | Midjourney and fal: accounts, submitting, and a worker that finishes jobs into the library |
| `models` | Model catalog from OmniGen listings, plus names, families and limits |
| `prompt_box` | The shared prompt box: editor, reference deck, keyframes, pickers, model selector |
| `pages` | Create Image, Video and Audio, Home, Library, and the shared page shell |
| `feed` | Generation feed: running, failed and finished jobs, library paging, grid and list |
| `overlays` | Lightbox, library picker, sign-in and settings dialogs |
| `overlays` | Lightbox, library picker, sign-in, settings, provider setup and the Midjourney sign-in window |
| `shell` | Custom title bar, sidebar flush to the window edge, content panel header |
| `theme`, `ui` | Tokens, Lucide-style icons, widgets, toasts, window chrome (drag, resize, caption buttons) |
@@ -52,6 +54,11 @@ How the backend works:
with lenient structs.
- **Identity headers.** Every request sends the desktop identity headers: `Origin` and
`User-Agent`.
- **Other providers.** Images can also run on Midjourney (your own account) or fal (your own API
key), through `artcraft_router` like the Tauri app. The accounts live in the same files
(`~/Artcraft/credentials`), so connecting in either app works in both. Those jobs are tracked in
`~/Artcraft/state/native_provider_jobs.json` and their results uploaded into your library, resuming
after a restart.
## Prompt box parity
@@ -72,7 +79,7 @@ How the backend works:
| Validation: starting frame, text-only banner, limits, uploads | Done |
| Clear all (confirms when references are attached) | Done |
| `@Character` mentions and the Characters modal (create/edit/delete) | Done |
| Provider choice (Midjourney direct, fal) | Not yet; everything runs via ArtCraft |
| Provider choice (Midjourney direct, fal), "Select Provider" | Done (images; video runs on ArtCraft) |
| In-app video playback | Done with FFmpeg on the PATH; otherwise opens in the system player |
| Audio reference playback (click to play/stop) | Done |
| Animated video previews in the feed (with a stills toggle) | Done |
@@ -93,8 +100,9 @@ Feed and lightbox:
## Diagnostics
Set `ARTCRAFT_SCREENSHOT=<file.png>` to save a screenshot once the window has settled.
`ARTCRAFT_PAGE` (`home`, `create-image`, `create-video`, `create-audio` or `library`) picks the page, and
`ARTCRAFT_SCREENSHOT_EXIT=1` closes the app afterwards.
`ARTCRAFT_PAGE` (`home`, `create-image`, `create-video`, `create-audio` or `library`) picks the page,
`ARTCRAFT_OPEN` (`settings` or `midjourney-login`) opens a dialog, and `ARTCRAFT_SCREENSHOT_EXIT=1`
closes the app afterwards. The screenshot shows egui only, not the sign-in webview.
## Tests
+92 -6
View File
@@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize};
use crate::backend::audio::AudioPlayer;
use crate::backend::media_cache::MediaCache;
use crate::backend::wire::{Prompt, SessionUser};
use crate::backend::providers::Provider;
use crate::backend::{Backend, Event, Modality, SHARE_URL_BASE};
use crate::feed::grid::{RatioCache, ViewMode};
use crate::feed::store::{FeedNotice, FeedStore};
@@ -20,6 +21,7 @@ use crate::overlays::account::{LoginAction, LoginDialog, SettingsAction, setting
use crate::overlays::characters::{CharactersAction, CharactersModal};
use crate::overlays::library_picker::{LibraryPicker, PickerAction};
use crate::overlays::lightbox::{Lightbox, LightboxAction, fit_inside};
use crate::overlays::providers::{AccountAction, AccountsView, LoginEvent, MidjourneyLogin, SetupAction, setup_dialog};
use crate::pages::common::{self, AppRequest, Env, FeedKey};
use crate::pages::create_audio::{AudioSettings, CreateAudioPage};
use crate::pages::create_image::{CreateImagePage, ImageSettings};
@@ -102,6 +104,11 @@ pub struct ArtcraftApp {
/// Characters being made: (job token, name).
pending_characters: Vec<(String, String)>,
settings_open: bool,
/// "Set up {provider}", after generating needed an account that isn't connected.
provider_setup: Option<Provider>,
midjourney_login: Option<MidjourneyLogin>,
/// The fal key being typed (setup dialog and Settings).
fal_key_input: String,
preview: Option<String>,
requests: Vec<AppRequest>,
screenshot: Option<ScreenshotRequest>,
@@ -125,7 +132,20 @@ impl ArtcraftApp {
backend.load_models(modality);
}
backend.refresh_session();
Self { cache, player, toasts: Toasts::default(), catalog: Catalog::default(), user: None, credits: None, page: saved.page, enter_to_generate: saved.enter_to_generate, view_mode: saved.view_mode, autoplay: saved.autoplay, sidebar_open: saved.sidebar_open, image: CreateImagePage::new(saved.image), video: CreateVideoPage::new(saved.video), audio: CreateAudioPage::new(saved.audio), library: FeedStore::default(), ratios: RatioCache::default(), prompts: HashMap::new(), prompt_records: HashMap::new(), requested_prompts: HashSet::new(), pending_recreate: None, batch_kinds: HashMap::new(), lightbox: None, picker: None, login: None, characters_modal: None, characters: Vec::new(), pending_characters: Vec::new(), settings_open: false, preview: None, requests: Vec::new(), screenshot: ScreenshotRequest::from_env(), next_poll: 0.0, next_credits: CREDITS_INTERVAL, models_retry_at: None, backend }
let mut app = Self { cache, player, toasts: Toasts::default(), catalog: Catalog::default(), user: None, credits: None, page: saved.page, enter_to_generate: saved.enter_to_generate, view_mode: saved.view_mode, autoplay: saved.autoplay, sidebar_open: saved.sidebar_open, image: CreateImagePage::new(saved.image), video: CreateVideoPage::new(saved.video), audio: CreateAudioPage::new(saved.audio), library: FeedStore::default(), ratios: RatioCache::default(), prompts: HashMap::new(), prompt_records: HashMap::new(), requested_prompts: HashSet::new(), pending_recreate: None, batch_kinds: HashMap::new(), lightbox: None, picker: None, login: None, characters_modal: None, characters: Vec::new(), pending_characters: Vec::new(), settings_open: false, provider_setup: None, midjourney_login: None, fal_key_input: String::new(), preview: None, requests: Vec::new(), screenshot: ScreenshotRequest::from_env(), next_poll: 0.0, next_credits: CREDITS_INTERVAL, models_retry_at: None, backend };
app.restore_provider_jobs();
// Diagnostics: `ARTCRAFT_OPEN=settings` or `midjourney-login` opens that dialog at launch.
match std::env::var("ARTCRAFT_OPEN").as_deref() {
Ok("settings") => app.settings_open = true,
Ok("midjourney-login") => app.account_action(AccountAction::ConnectMidjourney),
_ => {},
}
app
}
/// Pending cards for the Midjourney and fal jobs still running.
fn restore_provider_jobs(&mut self) {
self.image.feed.restore_local(&self.backend.provider_jobs(), &self.catalog);
}
// --- Events from the backend -------------------------------------------------------------
@@ -139,6 +159,7 @@ impl ArtcraftApp {
for feed in self.feeds_mut() {
feed.reset();
}
self.restore_provider_jobs();
self.next_poll = 0.0;
}
if self.user.is_some() {
@@ -295,6 +316,23 @@ impl ArtcraftApp {
self.next_poll = 0.0;
},
Event::CharacterChanged => self.backend.load_characters(),
Event::ProviderSetupNeeded(provider) => {
self.image.generating = false;
self.provider_setup = Some(provider);
},
Event::ProviderJobFinished { job_id, batch_token } => {
if self.image.feed.finish_local(&job_id) {
self.toasts.success("Image generation complete!");
}
// The batch brings the images into the feed.
self.batch_kinds.insert(job_id.clone(), MediaKind::Image);
self.backend.load_batch(job_id, batch_token);
},
Event::ProviderJobFailed { job_id, reason } => {
self.image.feed.fail_local(&job_id, &reason);
self.toasts.error(reason);
},
Event::ProviderAccountsChanged => {},
Event::Deleted(token) => {
for feed in self.feeds_mut() {
feed.remove_item(&token);
@@ -624,7 +662,7 @@ impl ArtcraftApp {
}
}
fn overlays(&mut self, ctx: &egui::Context) {
fn overlays(&mut self, ctx: &egui::Context, frame: &eframe::Frame) {
if let Some((kind, token)) = self.lightbox.as_ref().map(|lb| (lb.kind, lb.token.clone())) {
let key = if self.page == Page::Library { FeedKey::Library } else { FeedKey::of(kind) };
let order = self.feed(key).ordered_tokens();
@@ -706,9 +744,11 @@ impl ArtcraftApp {
}
}
}
if self.settings_open {
if self.settings_open && self.midjourney_login.is_none() {
let username = self.user.as_ref().map(|u| u.username.clone());
match settings_dialog(ctx, &mut self.enter_to_generate, username.as_deref(), self.credits, self.backend.data_root()) {
let accounts = self.backend.accounts();
let view = AccountsView { midjourney: accounts.midjourney_account(), fal_key_hint: accounts.fal_key_hint() };
match settings_dialog(ctx, &mut self.enter_to_generate, username.as_deref(), self.credits, &view, &mut self.fal_key_input, self.backend.data_root()) {
Some(SettingsAction::Close) => self.settings_open = false,
Some(SettingsAction::SignOut) => {
self.settings_open = false;
@@ -718,6 +758,32 @@ impl ArtcraftApp {
self.settings_open = false;
self.open_login();
},
Some(SettingsAction::Account(action)) => self.account_action(action),
None => {},
}
}
if let Some(provider) = self.provider_setup.filter(|_| self.midjourney_login.is_none()) {
match setup_dialog(ctx, provider, &mut self.fal_key_input) {
Some(SetupAction::Close) => self.provider_setup = None,
Some(SetupAction::ConnectMidjourney) => {
self.provider_setup = None;
self.account_action(AccountAction::ConnectMidjourney);
},
Some(SetupAction::SaveFalKey) => {
self.provider_setup = None;
self.account_action(AccountAction::SaveFalKey);
},
None => {},
}
}
if let Some(login) = &mut self.midjourney_login {
match login.show(ctx, frame) {
Some(LoginEvent::Captured(cookies)) => {
self.midjourney_login = None;
self.backend.save_midjourney_login(cookies);
self.toasts.success("Midjourney connected");
},
Some(LoginEvent::Closed) => self.midjourney_login = None,
None => {},
}
}
@@ -740,6 +806,26 @@ impl ArtcraftApp {
}
}
/// Connects or disconnects a provider account.
fn account_action(&mut self, action: AccountAction) {
match action {
AccountAction::ConnectMidjourney => self.midjourney_login = Some(MidjourneyLogin::new(self.backend.data_root())),
AccountAction::DisconnectMidjourney => self.backend.disconnect_midjourney(),
AccountAction::SaveFalKey => match self.backend.accounts().set_fal_key(&self.fal_key_input) {
Ok(()) => {
self.fal_key_input.clear();
self.toasts.success("FAL API key saved");
},
Err(err) => self.toasts.error(format!("Failed to save API key: {err}")),
},
AccountAction::RemoveFalKey => {
if let Err(err) = self.backend.accounts().clear_fal_key() {
self.toasts.error(format!("Failed to clear provider credential: {err}"));
}
},
}
}
fn lightbox_action(&mut self, ctx: &egui::Context, action: LightboxAction, item: &FeedItem) {
match action {
LightboxAction::Close => self.lightbox = None,
@@ -782,7 +868,7 @@ impl eframe::App for ArtcraftApp {
self.timers(ctx);
}
fn ui(&mut self, ui: &mut Ui, _frame: &mut eframe::Frame) {
fn ui(&mut self, ui: &mut Ui, frame: &mut eframe::Frame) {
let ctx = ui.ctx().clone();
self.cache.poll(&ctx);
if let Some(err) = self.player.poll() {
@@ -795,7 +881,7 @@ impl eframe::App for ArtcraftApp {
for request in std::mem::take(&mut self.requests) {
self.handle_request(&ctx, request);
}
self.overlays(&ctx);
self.overlays(&ctx, frame);
self.toasts.show(&ctx);
window::window_edges(&ctx);
if let Some(shot) = &mut self.screenshot {
+122 -1
View File
@@ -3,6 +3,7 @@
pub mod audio;
pub mod media_cache;
pub mod providers;
pub mod session;
pub mod video;
pub mod wire;
@@ -29,6 +30,7 @@ use artcraft_client::error::storyteller_error::StorytellerError;
use artcraft_client::utils::api_host::ApiHost;
use artcraft_client::utils::basic_json_post_request::basic_json_post_request;
use artcraft_client::utils::login_challenge_client::{LoginChallengeClient, LoginChallengeClientError};
use cookie_store_wrapper::cookie_store::CookieStore;
use enums::common::payments_namespace::PaymentsNamespace;
use log::{error, warn};
use serde::de::DeserializeOwned;
@@ -37,6 +39,11 @@ use tokens::tokens::media_files::MediaFileToken;
use crate::feed::types::MediaKind;
use crate::models::{self, ModelInfo};
use providers::Provider;
use providers::accounts::ProviderAccounts;
use providers::jobs::{JobStore, ProviderJob};
use providers::submit::SubmitError;
use providers::worker::Worker;
use session::Credentials;
use wire::{AudioModels, BatchMedia, CharactersPage, MediaFile, PromptResponse, SessionInfo, SessionJobs, SessionUser, UploadResponse, UserMediaList};
@@ -47,6 +54,8 @@ const MODEL_LIST_ATTEMPTS: u32 = 3;
const MODEL_LIST_BACKOFF: Duration = Duration::from_millis(250);
/// The share link base (`SHARE_URL_BASE`).
pub const SHARE_URL_BASE: &str = "https://getartcraft.com/media/";
/// fal's cost estimate (the provider-aware legacy endpoint; OmniGen only prices ArtCraft).
const PROVIDER_COST_PATH: &str = "/v1/generate/cost_estimate/image";
/// Characters are made for this model (the Tauri app's `CharactersModal`).
const CHARACTER_MODEL: &str = "seedance_2p0";
/// Sent on every request (`artcraft_client_identity::ARTCRAFT_DESKTOP_USER_AGENT`).
@@ -136,6 +145,19 @@ pub enum Event {
name: String,
},
CharacterChanged,
/// Generating with `Provider` needs its account connected first.
ProviderSetupNeeded(Provider),
/// A Midjourney or fal job's results are in the library, as `batch_token`.
ProviderJobFinished {
job_id: String,
batch_token: String,
},
ProviderJobFailed {
job_id: String,
reason: String,
},
/// A provider account was connected or disconnected.
ProviderAccountsChanged,
Deleted(String),
Saved(PathBuf),
Toast {
@@ -151,6 +173,8 @@ pub struct EnqueueMeta {
pub model_id: String,
pub batch_count: u32,
pub ref_image: Option<String>,
/// Who runs it (other providers' jobs are tracked locally).
pub provider: Provider,
}
/// A finished upload: the media token, plus what the server knows about the file.
@@ -172,6 +196,8 @@ pub struct Backend {
creds: Arc<RwLock<Option<Credentials>>>,
http: reqwest::Client,
device_login_cancel: Arc<AtomicBool>,
accounts: ProviderAccounts,
jobs: JobStore,
}
impl Backend {
@@ -182,7 +208,10 @@ impl Backend {
let host = session::api_host(&root);
let creds = session::load(&root);
let http = reqwest::Client::builder().user_agent(USER_AGENT).gzip(true).timeout(Duration::from_secs(60)).build().unwrap_or_default();
Self { rt, ctx, tx, rx, host, root, creds: Arc::new(RwLock::new(creds)), http, device_login_cancel: Arc::new(AtomicBool::new(false)) }
let (accounts, jobs) = (ProviderAccounts::load(&root), JobStore::load(&root));
let creds = Arc::new(RwLock::new(creds));
rt.spawn(Worker { host: host.clone(), creds: creds.clone(), http: http.clone(), accounts: accounts.clone(), jobs: jobs.clone(), tx: tx.clone(), ctx: ctx.clone() }.run());
Self { rt, ctx, tx, rx, host, root, creds, http, device_login_cancel: Arc::new(AtomicBool::new(false)), accounts, jobs }
}
pub fn runtime(&self) -> tokio::runtime::Handle {
@@ -408,6 +437,82 @@ impl Backend {
});
}
// --- Other providers ---------------------------------------------------------------------
/// Starts an image generation at Midjourney or fal; `image_urls` are the references' links.
pub fn generate_with_provider(&self, provider: Provider, fields: Map<String, Value>, image_urls: Vec<String>, meta: EnqueueMeta) {
let (api, accounts, jobs) = (self.api(), self.accounts.clone(), self.jobs.clone());
self.spawn(move |tx| async move {
match providers::submit::submit(&api.host, api.creds.as_ref(), &accounts, provider, &fields, image_urls).await {
Ok(job) => {
let id = job.id.clone();
jobs.put(job);
let _ = tx.send(Event::Enqueued { modality: Modality::Image, job_tokens: vec![id], meta });
},
Err(SubmitError::NeedsAccount(provider)) => {
let _ = tx.send(Event::ProviderSetupNeeded(provider));
},
Err(SubmitError::Message(message)) => {
let _ = tx.send(Event::EnqueueFailed { modality: Modality::Image, message, status: None });
},
}
});
}
/// What an image generation costs at `provider`: nothing extra for direct Midjourney (it uses
/// the user's subscription), fal's price from the server.
pub fn estimate_provider_cost(&self, provider: Provider, key: String, fields: Map<String, Value>) {
match provider {
Provider::Artcraft => self.estimate_cost(Modality::Image, key, fields),
Provider::Midjourney => {
let _ = self.tx.send(Event::Cost { key, credits: Some(0) });
},
Provider::Fal => {
let api = self.api();
self.spawn(move |tx| async move {
let credits = match basic_json_post_request::<_, Value>(&api.host, PROVIDER_COST_PATH, None, fal_cost_body(&fields)).await {
Ok(response) => response.get("cost_in_credits").and_then(Value::as_u64),
Err(err) => {
warn!("FAL cost estimate failed: {}", api_message(&err));
None
},
};
let _ = tx.send(Event::Cost { key, credits });
});
},
}
}
pub fn accounts(&self) -> &ProviderAccounts {
&self.accounts
}
/// Midjourney and fal jobs still running (they get pending cards at startup).
pub fn provider_jobs(&self) -> Vec<ProviderJob> {
self.jobs.all()
}
/// Keeps the Midjourney login captured by the sign-in window.
pub fn save_midjourney_login(&self, cookies: CookieStore) {
let accounts = self.accounts.clone();
self.spawn(move |tx| async move {
if let Err(err) = accounts.midjourney.save_login(cookies).await {
let _ = tx.send(Event::Toast { error: true, message: format!("Couldn't save the Midjourney login: {err}") });
}
let _ = tx.send(Event::ProviderAccountsChanged);
});
}
pub fn disconnect_midjourney(&self) {
let accounts = self.accounts.clone();
self.spawn(move |tx| async move {
if let Err(err) = accounts.midjourney.clear_credentials().await {
let _ = tx.send(Event::Toast { error: true, message: format!("Couldn't disconnect Midjourney: {err}") });
}
let _ = tx.send(Event::ProviderAccountsChanged);
});
}
// --- Uploads -----------------------------------------------------------------------------
/// Uploads a reference file from disk; images, MP4 videos and audio.
@@ -775,6 +880,15 @@ fn session_from_set_cookie(header: &str) -> Option<String> {
cookie::Cookie::parse(header).ok().filter(|c| c.name() == "session").map(|c| c.value().to_owned())
}
/// The fal cost request (`estimate_image_cost_command`): the page's fields, its references
/// counted as the generation mode.
fn fal_cost_body(fields: &Map<String, Value>) -> Value {
let images = fields.get("image_media_tokens").and_then(Value::as_array).map_or(0, Vec::len);
let mode = if images == 0 { serde_json::json!({ "type": "text_to_image" }) } else { serde_json::json!({ "type": "image_edit", "count": images }) };
let field = |name: &str| fields.get(name).cloned().unwrap_or(Value::Null);
serde_json::json!({ "model": field("model"), "provider": "fal", "generation_mode": mode, "aspect_ratio": field("aspect_ratio"), "resolution": field("resolution"), "quality": field("quality"), "image_batch_count": field("image_batch_count") })
}
/// All job tokens a generate call created (`inference_job_token` plus `all_job_tokens`).
fn job_tokens(response: &Value) -> Vec<String> {
let mut tokens: Vec<String> = response.get("inference_job_token").and_then(Value::as_str).map(str::to_owned).into_iter().collect();
@@ -849,6 +963,13 @@ mod tests {
assert_eq!(job_tokens(&serde_json::json!({ "inference_job_token": "only" })), ["only"]);
}
#[test]
fn fal_costs_count_references_as_the_mode() {
let fields = serde_json::json!({ "model": "flux_1_dev", "prompt": "x", "aspect_ratio": "square", "image_batch_count": 2, "image_media_tokens": ["mf_estimate_0"] });
let body = fal_cost_body(fields.as_object().unwrap());
assert_eq!(body, serde_json::json!({ "model": "flux_1_dev", "provider": "fal", "generation_mode": { "type": "image_edit", "count": 1 }, "aspect_ratio": "square", "resolution": null, "quality": null, "image_batch_count": 2 }));
}
#[test]
fn reads_the_servers_error_message() {
assert_eq!(body_message(r#"{"success":false,"error_message":"Not enough credits"}"#).as_deref(), Some("Not enough credits"));
@@ -0,0 +1,93 @@
//! The provider accounts, in the same files as the Tauri app (`~/Artcraft/credentials`), so an
//! account connected in either app works in both: the Midjourney login
//! (`midjourney_cookies.toml`) and the fal API key (`fal.api_key.txt`).
use std::path::{Path, PathBuf};
use std::sync::Arc;
use log::warn;
use midjourney_client::credentials::credential_manager::MidjourneyCredentialManager;
const MIDJOURNEY_FILENAME: &str = "midjourney_cookies.toml";
/// The Tauri app's pre-TOML Midjourney state, removed on disconnect so it can't come back.
const MIDJOURNEY_LEGACY_FILENAME: &str = "midjourney_state.json";
const FAL_KEY_FILENAME: &str = "fal.api_key.txt";
#[derive(Clone)]
pub struct ProviderAccounts {
pub midjourney: MidjourneyCredentialManager,
fal_key_path: Arc<PathBuf>,
}
impl ProviderAccounts {
pub fn load(data_root: &Path) -> Self {
let dir = data_root.join("credentials");
if let Err(err) = std::fs::create_dir_all(&dir) {
warn!("Couldn't create {}: {err}", dir.display());
}
let (path, legacy) = (dir.join(MIDJOURNEY_FILENAME), vec![dir.join(MIDJOURNEY_LEGACY_FILENAME)]);
let midjourney = MidjourneyCredentialManager::load(path.clone(), legacy.clone()).unwrap_or_else(|err| {
warn!("Couldn't load the Midjourney login: {err}");
MidjourneyCredentialManager::new(path, legacy, None)
});
Self { midjourney, fal_key_path: Arc::new(dir.join(FAL_KEY_FILENAME)) }
}
/// The connected Midjourney account's email ("" when it's connected but unknown), or `None`.
pub fn midjourney_account(&self) -> Option<String> {
match (self.midjourney.maybe_copy_cookie_store(), self.midjourney.maybe_account_info()) {
(Ok(None), Ok(None)) | (Err(_), _) | (_, Err(_)) => None,
(_, Ok(info)) => Some(info.and_then(|i| i.maybe_email).unwrap_or_default()),
}
}
pub fn fal_key(&self) -> Option<String> {
std::fs::read_to_string(&*self.fal_key_path).ok().map(|k| k.trim().to_owned()).filter(|k| !k.is_empty())
}
/// The key as Settings shows it (`first6********`).
pub fn fal_key_hint(&self) -> Option<String> {
self.fal_key().map(|k| mask_key(&k))
}
pub fn set_fal_key(&self, key: &str) -> std::io::Result<()> {
std::fs::write(&*self.fal_key_path, key.trim())
}
pub fn clear_fal_key(&self) -> std::io::Result<()> {
match std::fs::remove_file(&*self.fal_key_path) {
Err(err) if err.kind() != std::io::ErrorKind::NotFound => Err(err),
_ => Ok(()),
}
}
}
/// `first6********`, or just the stars for short keys (`provider_list_command`).
fn mask_key(key: &str) -> String {
if key.chars().count() > 6 {
format!("{}********", key.chars().take(6).collect::<String>())
} else {
"********".to_owned()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fal_keys_are_trimmed_masked_and_cleared() {
let dir = std::env::temp_dir().join(format!("artcraft-native-accounts-{}", uuid::Uuid::new_v4()));
let accounts = ProviderAccounts::load(&dir);
assert_eq!(accounts.fal_key(), None);
accounts.set_fal_key(" abcdef123456:secret \n").unwrap();
assert_eq!(accounts.fal_key().as_deref(), Some("abcdef123456:secret"));
assert_eq!(accounts.fal_key_hint().as_deref(), Some("abcdef********"));
accounts.clear_fal_key().unwrap();
accounts.clear_fal_key().unwrap();
assert_eq!(accounts.fal_key(), None);
assert_eq!(mask_key("short"), "********");
assert_eq!(accounts.midjourney_account(), None);
let _ = std::fs::remove_dir_all(dir);
}
}
@@ -0,0 +1,122 @@
//! Generations running at Midjourney or fal, saved (`state/native_provider_jobs.json`) so they
//! finish even after a restart. Each also checkpoints its upload into the user's library: the
//! batch it goes into, one idempotency token per result, and what's uploaded so far.
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use log::warn;
use serde::{Deserialize, Serialize};
use super::Provider;
const FILENAME: &str = "native_provider_jobs.json";
/// One generation at another provider.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ProviderJob {
/// Ours (`local_<uuid>`): the feed's pending card is keyed by it.
pub id: String,
pub provider: Provider,
pub model_id: String,
pub prompt: String,
/// Unix seconds.
pub created_at: i64,
pub batch_count: u32,
/// The provider's job: Midjourney's job id, or fal's request id.
pub provider_job_id: String,
/// fal's queue URLs.
pub maybe_status_url: Option<String>,
pub maybe_response_url: Option<String>,
/// The ArtCraft prompt the results are attributed to.
pub maybe_prompt_token: Option<String>,
/// The library batch the results go into.
pub batch_token: String,
/// Upload idempotency tokens, one per result (made as needed).
pub upload_tokens: Vec<String>,
/// Media tokens uploaded so far, in result order.
pub uploaded: Vec<String>,
}
impl ProviderJob {
/// The idempotency token for result `index`, making (and keeping) new ones as needed.
pub fn upload_token(&mut self, index: usize) -> String {
while self.upload_tokens.len() <= index {
self.upload_tokens.push(uuid::Uuid::new_v4().to_string());
}
self.upload_tokens[index].clone()
}
}
/// The saved jobs, shared by the submitter, the worker and the UI.
#[derive(Clone)]
pub struct JobStore {
path: Arc<PathBuf>,
jobs: Arc<Mutex<Vec<ProviderJob>>>,
}
impl JobStore {
pub fn load(data_root: &Path) -> Self {
let path = data_root.join("state").join(FILENAME);
let jobs = std::fs::read_to_string(&path).ok().and_then(|text| serde_json::from_str(&text).map_err(|err| warn!("Ignoring unreadable {}: {err}", path.display())).ok()).unwrap_or_default();
Self { path: Arc::new(path), jobs: Arc::new(Mutex::new(jobs)) }
}
pub fn all(&self) -> Vec<ProviderJob> {
self.jobs.lock().map(|j| j.clone()).unwrap_or_default()
}
/// Adds or replaces (by id) a job, then saves.
pub fn put(&self, job: ProviderJob) {
self.mutate(|jobs| match jobs.iter_mut().find(|j| j.id == job.id) {
Some(existing) => *existing = job,
None => jobs.push(job),
});
}
pub fn remove(&self, id: &str) {
self.mutate(|jobs| jobs.retain(|j| j.id != id));
}
fn mutate(&self, change: impl FnOnce(&mut Vec<ProviderJob>)) {
let Ok(mut jobs) = self.jobs.lock() else {
return;
};
change(&mut jobs);
if let Err(err) = save(&self.path, &jobs) {
warn!("Couldn't save {}: {err}", self.path.display());
}
}
}
/// Writes atomically, so a crash can't truncate the list.
fn save(path: &Path, jobs: &[ProviderJob]) -> anyhow::Result<()> {
let dir = path.parent().ok_or_else(|| anyhow::anyhow!("no state directory"))?;
std::fs::create_dir_all(dir)?;
let mut file = tempfile::NamedTempFile::new_in(dir)?;
file.write_all(serde_json::to_string_pretty(jobs)?.as_bytes())?;
file.persist(path)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn jobs_survive_a_restart_and_keep_their_upload_tokens() {
let dir = std::env::temp_dir().join(format!("artcraft-native-jobs-{}", uuid::Uuid::new_v4()));
let store = JobStore::load(&dir);
let mut job = ProviderJob { id: "local_1".into(), provider: Provider::Midjourney, model_id: "midjourney_7".into(), prompt: "a fox".into(), created_at: 1, batch_count: 4, provider_job_id: "mj".into(), maybe_status_url: None, maybe_response_url: None, maybe_prompt_token: None, batch_token: "batch_1".into(), upload_tokens: Vec::new(), uploaded: Vec::new() };
let first = job.upload_token(2);
assert_eq!(job.upload_tokens.len(), 3);
assert_eq!(job.upload_token(2), first, "a retried upload reuses its token");
store.put(job.clone());
let reloaded = JobStore::load(&dir);
assert_eq!(reloaded.all(), vec![job]);
reloaded.remove("local_1");
assert!(JobStore::load(&dir).all().is_empty());
let _ = std::fs::remove_dir_all(dir);
}
}
@@ -0,0 +1,110 @@
//! Generating through another provider (`ClassyModelSelector`'s "Select Provider"): Midjourney
//! with the user's own account, or fal with their API key, instead of ArtCraft. Images only;
//! video always runs on ArtCraft. Jobs are tracked here (they aren't ArtCraft jobs) and their
//! results uploaded into the user's library, like the Tauri app.
pub mod accounts;
pub mod jobs;
pub mod submit;
pub mod worker;
use serde::{Deserialize, Serialize};
use crate::models::ModelInfo;
/// Midjourney's models; direct jobs always make four images and take no references.
const MIDJOURNEY_MODELS: [&str; 4] = ["midjourney", "midjourney_7", "midjourney_7_niji", "midjourney_8"];
const MIDJOURNEY_BATCH: u16 = 4;
/// Who runs a generation.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Provider {
#[default]
Artcraft,
Midjourney,
Fal,
}
impl Provider {
/// The picker's name (`getProviderDisplayName`).
pub fn label(self) -> &'static str {
match self {
Provider::Artcraft => "ArtCraft",
Provider::Midjourney => "Midjourney",
Provider::Fal => "FAL",
}
}
/// The logo's key in the creator icons.
pub fn icon(self) -> &'static str {
match self {
Provider::Artcraft => "artcraft",
Provider::Midjourney => "midjourney",
Provider::Fal => "fal",
}
}
/// The providers `model_id` can run on, the default first (`ImageModels.ts`,
/// `buildModelsFromListing.ts`): fal for four ArtCraft models, Midjourney directly for
/// Midjourney's.
pub fn for_model(model_id: &str) -> &'static [Provider] {
match model_id {
"flux_1_dev" | "flux_1_schnell" | "nano_banana_2" | "nano_banana_pro" => &[Provider::Artcraft, Provider::Fal],
"midjourney" => &[Provider::Midjourney, Provider::Artcraft],
"midjourney_7" | "midjourney_7_niji" | "midjourney_8" => &[Provider::Artcraft, Provider::Midjourney],
_ => &[Provider::Artcraft],
}
}
/// `picked` if `model_id` runs there, else the model's default.
pub fn resolve(model_id: &str, picked: Option<Provider>) -> Provider {
let allowed = Provider::for_model(model_id);
picked.filter(|p| allowed.contains(p)).unwrap_or(allowed[0])
}
/// What the model can do on this provider (`imageModelForProvider`): direct Midjourney makes
/// four images, without references, resolution or quality choices.
pub fn adapt(self, model: &ModelInfo) -> ModelInfo {
let mut m = model.clone();
if self == Provider::Midjourney && is_midjourney(&model.id) {
m.batch_sizes = vec![MIDJOURNEY_BATCH];
m.batch_default = MIDJOURNEY_BATCH;
m.image_refs_max = 0;
m.resolutions.clear();
m.resolution_default = None;
m.qualities.clear();
m.quality_default = None;
}
m
}
}
pub fn is_midjourney(model_id: &str) -> bool {
MIDJOURNEY_MODELS.contains(&model_id)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn models_offer_their_providers_default_first() {
assert_eq!(Provider::for_model("nano_banana_pro"), [Provider::Artcraft, Provider::Fal]);
assert_eq!(Provider::for_model("midjourney"), [Provider::Midjourney, Provider::Artcraft]);
assert_eq!(Provider::for_model("midjourney_8"), [Provider::Artcraft, Provider::Midjourney]);
assert_eq!(Provider::for_model("seedream_4"), [Provider::Artcraft]);
assert_eq!(Provider::resolve("seedream_4", Some(Provider::Fal)), Provider::Artcraft, "a pick the model can't use falls back");
assert_eq!(Provider::resolve("flux_1_dev", Some(Provider::Fal)), Provider::Fal);
}
#[test]
fn direct_midjourney_makes_four_images_without_options() {
let model = ModelInfo { id: "midjourney_7".into(), batch_sizes: vec![1, 2, 4], batch_default: 1, image_refs_max: 1, resolutions: vec!["one_k".into()], aspect_ratios: vec!["square".into()], ..Default::default() };
let direct = Provider::Midjourney.adapt(&model);
assert_eq!((direct.batch_sizes.as_slice(), direct.batch_default, direct.image_refs_max), ([4].as_slice(), 4, 0));
assert!(direct.resolutions.is_empty());
assert_eq!(direct.aspect_ratios, model.aspect_ratios, "aspect ratios stay");
assert_eq!(Provider::Artcraft.adapt(&model).batch_sizes, model.batch_sizes);
}
}
@@ -0,0 +1,184 @@
//! Starting an image generation at Midjourney or fal (the Tauri app's `handle_midjourney` and
//! `handle_fal`): the ArtCraft prompt record, then the provider's request through the router.
use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
use artcraft_client::endpoints::prompts::create_prompt::create_prompt;
use artcraft_client::utils::api_host::ApiHost;
use artcraft_router::api::image_list_ref::ImageListRef;
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_fal_client::RouterFalClient;
use artcraft_router::client::router_midjourney_client::RouterMidjourneyClient;
use artcraft_router::errors::artcraft_router_error::ArtcraftRouterError;
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 artcraft_router::generate::generate_image::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
use log::warn;
use serde::de::DeserializeOwned;
use serde_json::{Map, Value};
use tokens::tokens::batch_generations::BatchGenerationToken;
use super::accounts::ProviderAccounts;
use super::jobs::ProviderJob;
use super::{Provider, is_midjourney};
use crate::backend::api_message;
/// Direct Midjourney jobs always make this many images.
const MIDJOURNEY_BATCH: u16 = 4;
/// Why a generation couldn't start.
#[derive(Debug, PartialEq)]
pub enum SubmitError {
/// The provider's account isn't connected (or its login expired).
NeedsAccount(Provider),
Message(String),
}
impl From<String> for SubmitError {
fn from(message: String) -> Self {
SubmitError::Message(message)
}
}
impl From<&str> for SubmitError {
fn from(message: &str) -> Self {
SubmitError::Message(message.to_owned())
}
}
/// Starts `fields` (the image page's OmniGen request) at `provider`; `image_urls` are the
/// references' CDN links.
pub async fn submit(host: &ApiHost, creds: Option<&StorytellerCredentialSet>, accounts: &ProviderAccounts, provider: Provider, fields: &Map<String, Value>, image_urls: Vec<String>) -> Result<ProviderJob, SubmitError> {
match provider {
Provider::Midjourney => submit_midjourney(host, creds, accounts, fields, image_urls).await,
Provider::Fal => submit_fal(host, creds, accounts, fields, image_urls).await,
Provider::Artcraft => Err("ArtCraft generations go through OmniGen".into()),
}
}
async fn submit_midjourney(host: &ApiHost, creds: Option<&StorytellerCredentialSet>, accounts: &ProviderAccounts, fields: &Map<String, Value>, image_urls: Vec<String>) -> Result<ProviderJob, SubmitError> {
let model_id = model_id(fields);
if !is_midjourney(model_id) {
return Err("This model cannot use a Midjourney account".into());
}
if !image_urls.is_empty() {
return Err("Direct Midjourney does not support image references yet".into());
}
if fields.get("image_batch_count").and_then(Value::as_u64).is_some_and(|n| n != u64::from(MIDJOURNEY_BATCH)) {
return Err("Direct Midjourney generates four images per job".into());
}
let mut builder = request_builder(RouterProvider::Midjourney, fields, Vec::new())?;
builder.resolution = None;
builder.quality = None;
builder.image_batch_count = Some(MIDJOURNEY_BATCH);
let prompt_request = router_image_request_to_artcraft_prompt(&builder);
let Ok(ImageGenerationDraftOrRequest::Request(generation)) = builder.build2() else {
return Err("Unexpected Midjourney request draft".into());
};
let session = accounts.midjourney.session().await.map_err(|err| err.to_string())?.ok_or(SubmitError::NeedsAccount(Provider::Midjourney))?;
// The prompt comes first, so the results can be attributed to it after a restart.
let creds = creds.ok_or("Please sign in to generate.")?;
let prompt = create_prompt(host, Some(creds), prompt_request).await.map_err(|err| api_message(&err))?;
let client = RouterClient::Midjourney(RouterMidjourneyClient::new(session.cookie_header, session.user_id, session.browser));
let GenerateImageResponse::Midjourney(response) = generation.send_request(&client).await.map_err(router_message(Provider::Midjourney))? else {
return Err("Unexpected Midjourney response".into());
};
Ok(new_job(Provider::Midjourney, fields, u32::from(MIDJOURNEY_BATCH), response.job_id, None, None, Some(prompt.prompt_token.to_string())))
}
async fn submit_fal(host: &ApiHost, creds: Option<&StorytellerCredentialSet>, accounts: &ProviderAccounts, fields: &Map<String, Value>, image_urls: Vec<String>) -> Result<ProviderJob, SubmitError> {
let key = accounts.fal_key().ok_or(SubmitError::NeedsAccount(Provider::Fal))?;
let builder = request_builder(RouterProvider::Fal, fields, image_urls)?;
// The prompt record is a nicety here: the generation goes ahead without it.
let maybe_prompt_token = match creds {
Some(creds) => create_prompt(host, Some(creds), router_image_request_to_artcraft_prompt(&builder)).await.map(|p| p.prompt_token.to_string()).map_err(|err| warn!("Couldn't create the prompt record: {}", api_message(&err))).ok(),
None => None,
};
let generation = match builder.build2() {
Ok(ImageGenerationDraftOrRequest::Request(generation)) => generation,
Ok(ImageGenerationDraftOrRequest::Draft(_)) => return Err("Fal should not be sending draft requests".into()),
Err(err) => return Err(err.to_string().into()),
};
let client = RouterClient::Fal(RouterFalClient::new_polling_only_from_raw_key(&key));
let GenerateImageResponse::Fal(response) = generation.send_request(&client).await.map_err(router_message(Provider::Fal))? else {
return Err("Unexpected FAL response".into());
};
let request_id = response.request_id.or(response.gateway_request_id).ok_or("FAL didn't return a request id")?;
let count = fields.get("image_batch_count").and_then(Value::as_u64).map_or(1, |n| n.max(1) as u32);
Ok(new_job(Provider::Fal, fields, count, request_id, response.maybe_status_url, response.maybe_response_url, maybe_prompt_token))
}
/// The router's request for the page's fields. The setting values are the router's own names
/// (`wide_sixteen_by_nine`, `one_k`, `high`).
fn request_builder(provider: RouterProvider, fields: &Map<String, Value>, image_urls: Vec<String>) -> Result<GenerateImageRequestBuilder, String> {
let model = match (parse::<RouterImageModel>(fields.get("model")), model_id(fields)) {
(Some(model), _) => model,
// The generic Midjourney model is the newest version.
(None, "midjourney") => RouterImageModel::Midjourney8,
(None, id) => return Err(format!("{id} isn't available on this provider")),
};
let mismatch = if matches!(provider, RouterProvider::Midjourney) { GenerationModeMismatchStrategy::AbortGeneration } else { GenerationModeMismatchStrategy::GenerateAnyway };
Ok(GenerateImageRequestBuilder { model, provider, prompt: fields.get("prompt").and_then(Value::as_str).map(str::to_owned), image_inputs: (!image_urls.is_empty()).then_some(ImageListRef::Urls(image_urls)), resolution: parse(fields.get("resolution")), aspect_ratio: parse(fields.get("aspect_ratio")), quality: parse(fields.get("quality")), image_batch_count: fields.get("image_batch_count").and_then(Value::as_u64).map(|n| n as u16), horizontal_angle: None, vertical_angle: None, zoom: None, request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, generation_mode_mismatch_strategy: Some(mismatch), idempotency_token: None })
}
fn new_job(provider: Provider, fields: &Map<String, Value>, batch_count: u32, provider_job_id: String, maybe_status_url: Option<String>, maybe_response_url: Option<String>, maybe_prompt_token: Option<String>) -> ProviderJob {
ProviderJob { id: format!("local_{}", uuid::Uuid::new_v4()), provider, model_id: model_id(fields).to_owned(), prompt: fields.get("prompt").and_then(Value::as_str).unwrap_or_default().to_owned(), created_at: chrono::Utc::now().timestamp(), batch_count, provider_job_id, maybe_status_url, maybe_response_url, maybe_prompt_token, batch_token: BatchGenerationToken::generate().to_string(), upload_tokens: Vec::new(), uploaded: Vec::new() }
}
/// What the user sees for a router error: billing problems point at the provider's website
/// (`GenericProviderBillingBlock`).
fn router_message(provider: Provider) -> impl Fn(ArtcraftRouterError) -> SubmitError {
move |err| match err {
ArtcraftRouterError::ProviderBillingError(_) => SubmitError::Message(format!("Please set up {} on their website to use it with ArtCraft.", provider.label())),
other => SubmitError::Message(other.to_string()),
}
}
fn model_id(fields: &Map<String, Value>) -> &str {
fields.get("model").and_then(Value::as_str).unwrap_or_default()
}
fn parse<T: DeserializeOwned>(value: Option<&Value>) -> Option<T> {
value.cloned().and_then(|v| serde_json::from_value(v).ok())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn page_fields_become_the_routers_request() {
let fields = json!({ "model": "nano_banana_pro", "prompt": "a fox", "aspect_ratio": "wide_sixteen_by_nine", "resolution": "two_k", "quality": "high", "image_batch_count": 2 });
let b = request_builder(RouterProvider::Fal, fields.as_object().unwrap(), vec!["https://cdn/x.png".into()]).unwrap();
assert_eq!(b.model, RouterImageModel::NanoBananaPro);
assert_eq!((b.prompt.as_deref(), b.image_batch_count), (Some("a fox"), Some(2)));
assert!(b.aspect_ratio.is_some() && b.resolution.is_some() && b.quality.is_some() && b.image_inputs.is_some());
let generic = request_builder(RouterProvider::Midjourney, json!({ "model": "midjourney" }).as_object().unwrap(), Vec::new()).unwrap();
assert_eq!(generic.model, RouterImageModel::Midjourney8);
assert!(request_builder(RouterProvider::Fal, json!({ "model": "seedance_2p0" }).as_object().unwrap(), Vec::new()).is_err());
}
#[tokio::test]
async fn direct_midjourney_checks_its_limits_before_anything_else() {
let dir = std::env::temp_dir().join(format!("artcraft-native-submit-{}", uuid::Uuid::new_v4()));
let accounts = ProviderAccounts::load(&dir);
let host = ApiHost::Storyteller;
let run = |fields: Value, urls: Vec<String>| {
let (host, accounts) = (&host, &accounts);
async move { submit(host, None, accounts, Provider::Midjourney, fields.as_object().unwrap(), urls).await }
};
assert_eq!(run(json!({ "model": "flux_1_dev" }), Vec::new()).await.unwrap_err(), SubmitError::Message("This model cannot use a Midjourney account".into()));
assert_eq!(run(json!({ "model": "midjourney_7" }), vec!["u".into()]).await.unwrap_err(), SubmitError::Message("Direct Midjourney does not support image references yet".into()));
assert_eq!(run(json!({ "model": "midjourney_7", "image_batch_count": 2 }), Vec::new()).await.unwrap_err(), SubmitError::Message("Direct Midjourney generates four images per job".into()));
assert_eq!(run(json!({ "model": "midjourney_7", "prompt": "a fox", "image_batch_count": 4 }), Vec::new()).await.unwrap_err(), SubmitError::NeedsAccount(Provider::Midjourney), "no login yet");
let fal = submit(&host, None, &accounts, Provider::Fal, json!({ "model": "flux_1_dev" }).as_object().unwrap(), Vec::new()).await;
assert_eq!(fal.unwrap_err(), SubmitError::NeedsAccount(Provider::Fal));
let _ = std::fs::remove_dir_all(dir);
}
}
@@ -0,0 +1,335 @@
//! Finishing provider jobs in the background (the Tauri app's `midjourney_long_polling_thread`
//! and `poll_fal_tasks`): fal's queue is polled; Midjourney's feed is read every 15 s and its
//! websocket reports completions sooner. Finished results are downloaded and uploaded into the
//! user's library as one batch, resuming where they left off after a failure or restart.
use std::collections::HashSet;
use std::sync::mpsc::Sender;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
use artcraft_client::endpoints::media_files::upload_image_media_file_from_file::{UploadImageFromFileArgs, upload_image_media_file_from_file_with_idempotency_token};
use artcraft_client::utils::api_host::ApiHost;
use enums::common::generation_provider::GenerationProvider;
use fal_client::creds::fal_api_key::FalApiKey;
use fal_client::polling::poll_job_response::poll_job_response::{PollJobResponseArgs, poll_job_response};
use fal_client::polling::poll_job_response::success_case_extractors::PollResponseExtractedContents;
use fal_client::polling::poll_job_status::poll_job_status::{FalJobStatus, PollJobStatusArgs, poll_job_status};
use log::{info, warn};
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::{OpenMidjourneyWebSocketRequest, open_midjourney_websocket};
use midjourney_client::credentials::credential_manager::MidjourneySession;
use midjourney_client::endpoints::imagine::{ImagineArgs, ImagineRequest, imagine};
use midjourney_client::recipes::get_user_info::{GetUserInfoArgs, get_user_info};
use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
use tokens::tokens::batch_generations::BatchGenerationToken;
use tokens::tokens::prompts::PromptToken;
use tokio::sync::broadcast::Receiver;
use tokio::sync::broadcast::error::TryRecvError;
use tokio::time::timeout;
use super::Provider;
use super::accounts::ProviderAccounts;
use super::jobs::{JobStore, ProviderJob};
use crate::backend::Event;
use crate::backend::session::Credentials;
/// How often jobs are checked (fal's poll interval).
const TICK: Duration = Duration::from_secs(2);
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);
/// Each download and upload of a result.
const TRANSFER_TIMEOUT: Duration = Duration::from_secs(60);
/// A Midjourney job not seen finished by then is given up on.
const MIDJOURNEY_GIVE_UP_SECS: i64 = 60 * 60;
/// Direct Midjourney jobs make four images.
const MIDJOURNEY_RESULTS: usize = 4;
pub struct Worker {
pub host: ApiHost,
pub creds: Arc<RwLock<Option<Credentials>>>,
pub http: reqwest::Client,
pub accounts: ProviderAccounts,
pub jobs: JobStore,
pub tx: Sender<Event>,
pub ctx: egui::Context,
}
/// Where a job's results come from.
enum Results<'a> {
Urls(Vec<String>),
Midjourney(&'a ImageDownloaderClient),
}
impl Results<'_> {
fn count(&self) -> usize {
match self {
Results::Urls(urls) => urls.len(),
Results::Midjourney(_) => MIDJOURNEY_RESULTS,
}
}
}
/// What the Midjourney side remembers between ticks.
#[derive(Default)]
struct MidjourneyWatch {
socket: Option<Socket>,
/// Jobs seen finished (by the websocket or the feed).
finished: HashSet<String>,
next_feed: Option<Instant>,
next_socket_attempt: Option<Instant>,
}
struct Socket {
credential_id: String,
socket: MidjourneyWebSocket,
events: Receiver<Arc<MidjourneyWsEvent>>,
subscribed: HashSet<String>,
}
impl Worker {
pub async fn run(self) {
let mut watch = MidjourneyWatch::default();
loop {
self.tick(&mut watch).await;
tokio::time::sleep(TICK).await;
}
}
async fn tick(&self, watch: &mut MidjourneyWatch) {
let jobs = self.jobs.all();
let (midjourney, fal): (Vec<ProviderJob>, Vec<ProviderJob>) = jobs.into_iter().partition(|j| j.provider == Provider::Midjourney);
if midjourney.is_empty() {
watch.close();
}
// Results go into the library: nothing can finish until someone's signed in.
let Some(creds) = self.creds.read().ok().and_then(|c| c.as_ref().map(Credentials::to_set)) else {
return;
};
if !fal.is_empty() {
if let Some(key) = self.accounts.fal_key().map(|k| FalApiKey::from_str(&k)) {
for job in fal {
self.poll_fal(job, &key, &creds).await;
}
}
}
if !midjourney.is_empty() {
self.poll_midjourney(watch, midjourney, &creds).await;
}
}
async fn poll_fal(&self, job: ProviderJob, key: &FalApiKey, creds: &StorytellerCredentialSet) {
let Some(status_url) = job.maybe_status_url.clone() else {
return self.fail(&job, "No queue status URL to check");
};
let status = match timeout(REQUEST_TIMEOUT, poll_job_status(PollJobStatusArgs { status_url: &status_url, api_key: key })).await {
Ok(Ok(status)) => status,
other => return warn!("Polling FAL job {} failed: {other:?}", job.provider_job_id),
};
if status.status != FalJobStatus::Completed {
return;
}
let Some(response_url) = job.maybe_response_url.clone().or(status.response_url) else {
return self.fail(&job, "Job completed but no response URL available");
};
let response = match timeout(REQUEST_TIMEOUT, poll_job_response(PollJobResponseArgs { response_url: &response_url, api_key: key })).await {
Ok(Ok(response)) => response,
other => return warn!("Fetching FAL results for {} failed (will retry): {other:?}", job.provider_job_id),
};
let urls = image_urls(response.extracted_contents.as_ref());
if urls.is_empty() {
return self.fail(&job, "FAL returned no images");
}
if let Err(err) = self.finish(job, Results::Urls(urls), creds).await {
warn!("Uploading FAL results failed (will retry): {err}");
}
}
async fn poll_midjourney(&self, watch: &mut MidjourneyWatch, jobs: Vec<ProviderJob>, creds: &StorytellerCredentialSet) {
let session = match self.accounts.midjourney.session().await {
Ok(Some(session)) => session,
Ok(None) => return watch.close(),
Err(err) => return warn!("Midjourney login check failed: {err}"),
};
let now = chrono::Utc::now().timestamp();
for job in jobs.iter().filter(|j| now - j.created_at > MIDJOURNEY_GIVE_UP_SECS && !watch.finished.contains(&j.provider_job_id)) {
self.fail(job, "Midjourney didn't finish this job");
}
let pending: HashSet<String> = jobs.iter().filter(|j| now - j.created_at <= MIDJOURNEY_GIVE_UP_SECS || watch.finished.contains(&j.provider_job_id)).map(|j| j.provider_job_id.clone()).collect();
watch.finished.retain(|id| pending.contains(id));
watch.refresh(&session, &pending).await;
let downloader = match ImageDownloaderClient::create(Some(session.browser.clone())) {
Ok(downloader) => downloader,
Err(err) => return warn!("Midjourney downloader unavailable: {err:?}"),
};
let done: Vec<ProviderJob> = jobs.into_iter().filter(|j| watch.finished.contains(&j.provider_job_id)).collect();
for job in done {
let id = job.provider_job_id.clone();
match self.finish(job, Results::Midjourney(&downloader), creds).await {
Ok(()) => {
watch.finished.remove(&id);
},
Err(err) => warn!("Uploading Midjourney results for {id} failed (will retry): {err}"),
}
}
}
/// Uploads the results not uploaded yet, then reports the job done.
async fn finish(&self, mut job: ProviderJob, results: Results<'_>, creds: &StorytellerCredentialSet) -> anyhow::Result<()> {
let provider = match job.provider {
Provider::Midjourney => GenerationProvider::Midjourney,
Provider::Fal => GenerationProvider::Fal,
Provider::Artcraft => GenerationProvider::Artcraft,
};
let prompt = job.maybe_prompt_token.as_deref().map(PromptToken::new_from_str);
let batch = BatchGenerationToken::new_from_str(&job.batch_token);
for index in job.uploaded.len()..results.count() {
let (bytes, ext) = timeout(TRANSFER_TIMEOUT, self.download(&results, &job, index)).await??;
let mut file = tempfile::Builder::new().prefix("artcraft-provider-").suffix(&format!(".{ext}")).tempfile()?;
std::io::Write::write_all(&mut file, &bytes)?;
// The token is saved before the upload, so a retry can't upload the same result twice.
let token = job.upload_token(index);
self.jobs.put(job.clone());
let args = UploadImageFromFileArgs { api_host: &self.host, maybe_creds: Some(creds), path: file.path(), is_intermediate_system_file: false, maybe_prompt_token: prompt.as_ref(), maybe_generation_provider: Some(provider), maybe_batch_token: Some(&batch) };
let uploaded = timeout(TRANSFER_TIMEOUT, upload_image_media_file_from_file_with_idempotency_token(args, &token)).await??;
job.uploaded.push(uploaded.media_file_token.to_string());
self.jobs.put(job.clone());
}
info!("{:?} job {} finished with {} results", job.provider, job.provider_job_id, job.uploaded.len());
self.jobs.remove(&job.id);
self.send(Event::ProviderJobFinished { job_id: job.id, batch_token: job.batch_token });
Ok(())
}
async fn download(&self, results: &Results<'_>, job: &ProviderJob, index: usize) -> anyhow::Result<(Vec<u8>, &'static str)> {
match results {
Results::Urls(urls) => {
let url = &urls[index];
let bytes = self.http.get(url).send().await?.error_for_status()?.bytes().await?;
Ok((bytes.to_vec(), image_extension(url)))
},
Results::Midjourney(downloader) => Ok((downloader.download_image(&job.provider_job_id, index as u8).await?, "png")),
}
}
fn fail(&self, job: &ProviderJob, reason: &str) {
warn!("{:?} job {} failed: {reason}", job.provider, job.provider_job_id);
self.jobs.remove(&job.id);
self.send(Event::ProviderJobFailed { job_id: job.id.clone(), reason: reason.to_owned() });
}
fn send(&self, event: Event) {
let _ = self.tx.send(event);
self.ctx.request_repaint();
}
}
impl MidjourneyWatch {
/// Notes which pending jobs have finished: websocket events as they come, and the feed every
/// so often (it also covers frames the socket missed, and jobs from before a restart).
async fn refresh(&mut self, session: &MidjourneySession, pending: &HashSet<String>) {
if self.socket.as_ref().is_some_and(|s| s.credential_id != session.credential_id || !s.socket.is_connected()) {
self.close();
}
self.drain(pending);
let now = Instant::now();
if self.next_feed.is_none_or(|t| now >= t) {
self.next_feed = Some(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(feed)) => self.finished.extend(feed.items.into_iter().filter_map(|item| item.id).filter(|id| pending.contains(id))),
other => warn!("Reading the Midjourney feed failed: {other:?}"),
}
}
if self.socket.is_none() && self.next_socket_attempt.is_none_or(|t| now >= t) {
self.next_socket_attempt = Some(now + SOCKET_RETRY_INTERVAL);
match timeout(REQUEST_TIMEOUT, open_socket(session)).await {
Ok(Ok(socket)) => self.socket = Some(socket),
_ => log::debug!("Midjourney websocket unavailable; the feed still reports completions"),
}
}
if let Some(socket) = self.socket.as_mut() {
socket.subscribed.retain(|id| pending.contains(id));
for id in pending {
if !socket.subscribed.contains(id) && socket.socket.subscribe_to_job(id).is_ok() {
socket.subscribed.insert(id.clone());
}
}
}
}
fn drain(&mut self, pending: &HashSet<String>) {
let Some(socket) = self.socket.as_mut() else {
return;
};
loop {
match socket.events.try_recv() {
Ok(event) => {
if let MidjourneyWsEvent::Completed { job_id, .. } = event.as_ref() {
if pending.contains(job_id) {
self.finished.insert(job_id.clone());
}
}
},
// Missed frames: read the feed now.
Err(TryRecvError::Lagged(_)) => self.next_feed = None,
Err(TryRecvError::Empty | TryRecvError::Closed) => break,
}
}
}
fn close(&mut self) {
if let Some(socket) = self.socket.take() {
socket.socket.close();
}
}
}
async fn open_socket(session: &MidjourneySession) -> anyhow::Result<Socket> {
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(Socket { credential_id: session.credential_id.clone(), socket, events, subscribed: HashSet::new() })
}
/// A fal job's image links: the batch, else the single image.
fn image_urls(extracted: Option<&PollResponseExtractedContents>) -> Vec<String> {
let Some(extracted) = extracted else {
return Vec::new();
};
let batch: Vec<String> = extracted.images.iter().flatten().filter_map(|i| i.url.clone()).collect();
if !batch.is_empty() {
return batch;
}
extracted.image.as_ref().and_then(|i| i.url.clone()).into_iter().collect()
}
/// The file extension a result URL ends in (PNG when it doesn't say).
fn image_extension(url: &str) -> &'static str {
let path = url.split(['?', '#']).next().unwrap_or_default().to_ascii_lowercase();
["jpg", "jpeg", "webp", "png"].into_iter().find(|ext| path.ends_with(&format!(".{ext}"))).unwrap_or("png")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn results_keep_their_file_type() {
assert_eq!(image_extension("https://fal.media/files/a/b.jpeg?x=1"), "jpeg");
assert_eq!(image_extension("https://fal.media/files/a/b.webp"), "webp");
assert_eq!(image_extension("https://fal.media/files/a/b"), "png");
}
#[test]
fn fal_results_prefer_the_batch() {
let contents = fal_client::polling::poll_job_response::success_case_extractors::extract_contents_from_response(&serde_json::json!({ "images": [{ "url": "https://fal/1.png" }, { "url": "https://fal/2.png" }] }));
assert_eq!(image_urls(contents.as_ref()), ["https://fal/1.png", "https://fal/2.png"]);
let single = fal_client::polling::poll_job_response::success_case_extractors::extract_contents_from_response(&serde_json::json!({ "image": { "url": "https://fal/one.png" } }));
assert_eq!(image_urls(single.as_ref()), ["https://fal/one.png"]);
assert!(image_urls(None).is_empty());
}
}
+58 -7
View File
@@ -3,8 +3,10 @@
use std::collections::{HashMap, HashSet};
use crate::backend::wire::{Job, JobPhase, MediaFile, unix_secs};
use crate::backend::EnqueueMeta;
use crate::backend::providers::Provider;
use crate::backend::providers::jobs::ProviderJob;
use crate::backend::wire::{Job, JobPhase, MediaFile, unix_secs};
use crate::feed::types::{FailedJob, FeedItem, MediaKind, PendingJob};
use crate::models::{self, Catalog};
@@ -61,7 +63,7 @@ impl FeedStore {
self.meta.insert(token.clone(), meta.clone());
self.seen_running.insert(token.clone());
if !self.pending.iter().any(|p| &p.job_token == token) {
self.pending.insert(0, PendingJob { job_token: token.clone(), kind, prompt: meta.prompt.clone(), model_id: Some(meta.model_id.clone()), created_at: now, server_progress: None, expected_secs: expected, batch_count: meta.batch_count });
self.pending.insert(0, PendingJob { job_token: token.clone(), kind, prompt: meta.prompt.clone(), model_id: Some(meta.model_id.clone()), created_at: now, server_progress: None, expected_secs: expected, batch_count: meta.batch_count, local: meta.provider != Provider::Artcraft });
}
}
}
@@ -82,7 +84,7 @@ impl FeedStore {
JobPhase::Running => {
self.seen_running.insert(job.job_token.clone());
let expected = model_id.as_deref().and_then(|id| catalog.find(id)).map_or_else(|| models::expected_secs(model_id.as_deref().unwrap_or_default(), kind), |m| m.expected_secs);
pending.push(PendingJob { job_token: job.job_token.clone(), kind, prompt, model_id, created_at: unix_secs(&job.created_at), server_progress: Some(job.status.progress_percentage), expected_secs: expected, batch_count: meta.map_or(1, |m| m.batch_count) });
pending.push(PendingJob { job_token: job.job_token.clone(), kind, prompt, model_id, created_at: unix_secs(&job.created_at), server_progress: Some(job.status.progress_percentage), expected_secs: expected, batch_count: meta.map_or(1, |m| m.batch_count), local: false });
},
JobPhase::Failed => {
let updated = unix_secs(&job.updated_at);
@@ -90,7 +92,7 @@ impl FeedStore {
notices.push(FeedNotice::Failed(job.status.maybe_failure_message.clone().unwrap_or_else(|| job.status.failure_reason())));
}
if now - updated <= FAILURE_WINDOW_SECS {
failed.push(FailedJob { job_token: job.job_token.clone(), prompt, model_id, updated_at: updated, reason: job.status.failure_reason(), message: job.status.maybe_failure_message.clone().filter(|_| job.status.maybe_failure_category_updated.as_deref() != Some("unknown")), ref_image: meta.and_then(|m| m.ref_image.clone()) });
failed.push(FailedJob { job_token: job.job_token.clone(), prompt, model_id, updated_at: updated, reason: job.status.failure_reason(), message: job.status.maybe_failure_message.clone().filter(|_| job.status.maybe_failure_category_updated.as_deref() != Some("unknown")), ref_image: meta.and_then(|m| m.ref_image.clone()), local: false });
}
},
JobPhase::Succeeded => {
@@ -113,13 +115,15 @@ impl FeedStore {
},
}
}
// Jobs started a moment ago may not be in the list yet: keep them.
// Jobs started a moment ago may not be in the list yet, and other providers' never are: keep
// them.
for p in self.pending.drain(..) {
let listed = jobs.iter().any(|j| j.job_token == p.job_token);
if !listed && now - p.created_at < 60 && !pending.iter().any(|q: &PendingJob| q.job_token == p.job_token) {
if (p.local || (!listed && now - p.created_at < 60)) && !pending.iter().any(|q: &PendingJob| q.job_token == p.job_token) {
pending.push(p);
}
}
failed.extend(self.failed.drain(..).filter(|f| f.local));
// Everything finished before the first poll comes from the library instead.
self.first_poll_done = true;
pending.sort_by_key(|p| std::cmp::Reverse(p.created_at));
@@ -140,6 +144,35 @@ impl FeedStore {
}
}
/// Pending cards for the Midjourney and fal jobs still running (after a restart).
pub fn restore_local(&mut self, jobs: &[ProviderJob], catalog: &Catalog) {
for job in jobs {
if !self.pending.iter().any(|p| p.job_token == job.id) {
let expected = catalog.find(&job.model_id).map_or_else(|| models::expected_secs(&job.model_id, MediaKind::Image), |m| m.expected_secs);
self.pending.push(PendingJob { job_token: job.id.clone(), kind: MediaKind::Image, prompt: job.prompt.clone(), model_id: Some(job.model_id.clone()), created_at: job.created_at, server_progress: None, expected_secs: expected, batch_count: job.batch_count, local: true });
}
}
self.pending.sort_by_key(|p| std::cmp::Reverse(p.created_at));
}
/// A Midjourney or fal job's results are in: its card goes (the batch brings the images).
/// Returns whether it was this feed's.
pub fn finish_local(&mut self, job_id: &str) -> bool {
let before = self.pending.len();
self.pending.retain(|p| p.job_token != job_id);
self.pending.len() != before
}
/// A Midjourney or fal job failed: its card shows the reason until dismissed.
pub fn fail_local(&mut self, job_id: &str, reason: &str) -> bool {
let Some(i) = self.pending.iter().position(|p| p.job_token == job_id) else {
return false;
};
let p = self.pending.remove(i);
self.failed.insert(0, FailedJob { job_token: p.job_token, prompt: p.prompt, model_id: p.model_id, updated_at: chrono::Utc::now().timestamp(), reason: "Generation failed".to_owned(), message: Some(reason.to_owned()), ref_image: self.meta.get(job_id).and_then(|m| m.ref_image.clone()), local: true });
true
}
pub fn library_started(&self) -> bool {
self.library_started
}
@@ -254,6 +287,7 @@ fn job_kind(job: &Job, catalog: &Catalog) -> Option<MediaKind> {
#[cfg(test)]
mod tests {
use super::*;
use crate::backend::providers::Provider;
use crate::models::ModelInfo;
fn catalog() -> Catalog {
@@ -326,10 +360,27 @@ mod tests {
assert_eq!(store.pending.len(), 1);
}
#[test]
fn provider_jobs_outlive_polls_until_they_finish_or_fail() {
let (c, mut store) = (catalog(), FeedStore::default());
let meta = EnqueueMeta { prompt: "p".into(), model_id: "midjourney_7".into(), batch_count: 4, ref_image: None, provider: Provider::Midjourney };
store.add_enqueued(MediaKind::Image, &["local_a".into(), "local_b".into()], &meta, &c);
store.pending.iter_mut().for_each(|p| p.created_at -= 3600);
store.apply_jobs(MediaKind::Image, &[], &c);
assert_eq!(store.pending.len(), 2, "the server never lists them");
assert!(store.finish_local("local_a"));
assert!(store.fail_local("local_b", "Midjourney didn't finish this job"));
store.apply_jobs(MediaKind::Image, &[], &c);
assert!(store.pending.is_empty());
assert_eq!(store.failed[0].message.as_deref(), Some("Midjourney didn't finish this job"), "the failure stays until dismissed");
store.dismiss("local_b");
assert!(store.failed.is_empty());
}
#[test]
fn just_enqueued_jobs_survive_a_poll_that_predates_them() {
let (c, mut store) = (catalog(), FeedStore::default());
let meta = EnqueueMeta { prompt: "p".into(), model_id: "nano_banana_pro".into(), batch_count: 2, ref_image: None };
let meta = EnqueueMeta { prompt: "p".into(), model_id: "nano_banana_pro".into(), batch_count: 2, ref_image: None, provider: Provider::Artcraft };
store.add_enqueued(MediaKind::Image, &["fresh".into()], &meta, &c);
store.apply_jobs(MediaKind::Image, &[], &c);
assert_eq!(store.pending.len(), 1);
+4
View File
@@ -62,6 +62,8 @@ pub struct PendingJob {
/// How long this model usually takes, for the estimated progress bar.
pub expected_secs: f32,
pub batch_count: u32,
/// Running at Midjourney or fal (tracked here, not in the server's job list).
pub local: bool,
}
/// A generation that failed or was cancelled.
@@ -75,6 +77,8 @@ pub struct FailedJob {
pub message: Option<String>,
/// The first reference image, shown faintly behind the error.
pub ref_image: Option<String>,
/// A Midjourney or fal job (tracked here, not in the server's job list).
pub local: bool,
}
/// The progress and time-left labels for a pending job (`derivePendingStatus`).
+12 -3
View File
@@ -2,6 +2,7 @@
use egui::{Color32, Id, Sense, Stroke, Ui, pos2, vec2};
use crate::overlays::providers::{AccountAction, AccountsView, accounts_section};
use crate::shell::{WEB_APP_URL, group_thousands};
use crate::theme;
use crate::ui::icons::Icon;
@@ -105,10 +106,12 @@ pub enum SettingsAction {
Close,
SignOut,
SignIn,
Account(AccountAction),
}
/// The settings dialog (Settings → Misc and Account in the webapp).
pub fn settings_dialog(ctx: &egui::Context, enter_to_generate: &mut bool, username: Option<&str>, credits: Option<u64>, data_root: &std::path::Path) -> Option<SettingsAction> {
/// The settings dialog (Settings → Misc and Accounts in the webapp). `fal_key` is the key being
/// typed.
pub fn settings_dialog(ctx: &egui::Context, enter_to_generate: &mut bool, username: Option<&str>, credits: Option<u64>, accounts: &AccountsView, fal_key: &mut String, data_root: &std::path::Path) -> Option<SettingsAction> {
let mut action = None;
let close = widgets::modal(ctx, Id::new("settings"), 520.0, |ui| {
ui.horizontal(|ui| {
@@ -160,6 +163,12 @@ pub fn settings_dialog(ctx: &egui::Context, enter_to_generate: &mut bool, userna
},
}
ui.add_space(18.0);
widgets::hud(ui, "Provider accounts", theme::FAINT);
ui.add_space(6.0);
if let Some(account) = accounts_section(ui, accounts, fal_key) {
action = Some(SettingsAction::Account(account));
}
ui.add_space(18.0);
widgets::hud(ui, "About", theme::FAINT);
ui.add_space(6.0);
ui.label(egui::RichText::new(format!("ArtCraft native {}", env!("CARGO_PKG_VERSION"))).size(13.0).color(theme::MUTED));
@@ -176,7 +185,7 @@ pub fn settings_dialog(ctx: &egui::Context, enter_to_generate: &mut bool, userna
}
/// A single-line input on the sunken surface with a hairline (and a blue line while focused).
fn text_field(ui: &mut Ui, value: &mut String, password: bool, hint: &str) -> egui::Response {
pub(crate) fn text_field(ui: &mut Ui, value: &mut String, password: bool, hint: &str) -> egui::Response {
let id = ui.next_auto_id();
let focused = ui.memory(|m| m.has_focus(id));
egui::Frame::new().fill(theme::SUNKEN).stroke(Stroke::new(1.0, if focused { theme::ACCENT } else { theme::LINE })).corner_radius(theme::RADIUS).inner_margin(egui::Margin::symmetric(10, 8)).show(ui, |ui| ui.add(egui::TextEdit::singleline(value).id(id).password(password).frame(egui::Frame::NONE).desired_width(f32::INFINITY).hint_text(egui::RichText::new(hint).color(theme::FAINT)))).inner
+1
View File
@@ -2,3 +2,4 @@ pub mod account;
pub mod characters;
pub mod library_picker;
pub mod lightbox;
pub mod providers;
+261
View File
@@ -0,0 +1,261 @@
//! Connecting provider accounts: the "Set up {provider}" dialog (`ProviderSetupModal`), the
//! Midjourney sign-in window (a system webview, used only to capture the login cookies, like
//! the Tauri app's), and the accounts rows in Settings.
use std::path::Path;
use cookie_store_wrapper::cookie_store::CookieStore;
use egui::{Id, Rect, Sense, Ui, vec2};
use log::warn;
use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
use midjourney_client::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
use url::Url;
use wry::raw_window_handle::HasWindowHandle;
use wry::{WebContext, WebView, WebViewBuilder};
use crate::backend::providers::Provider;
use crate::overlays::account::text_field;
use crate::theme;
use crate::ui::icons::{self, Icon};
use crate::ui::widgets;
const MIDJOURNEY_URL: &str = "https://www.midjourney.com/";
/// The login cookies live on both hosts.
const MIDJOURNEY_ORIGINS: [&str; 2] = ["https://www.midjourney.com/", "https://midjourney.com/"];
/// How often the sign-in window is checked for the login cookies.
const COOKIE_CHECK_SECS: f64 = 2.0;
const LOGIN_WIDTH: f32 = 560.0;
/// What the setup dialog asks for.
#[derive(Debug, PartialEq)]
pub enum SetupAction {
Close,
ConnectMidjourney,
SaveFalKey,
}
/// "Set up {provider}": Midjourney signs in through its window; fal takes an API key here (the
/// Tauri app has no place for it yet).
pub fn setup_dialog(ctx: &egui::Context, provider: Provider, fal_key: &mut String) -> Option<SetupAction> {
let name = provider.label();
let mut action = None;
let close = widgets::modal(ctx, Id::new("provider-setup"), 460.0, |ui| {
ui.horizontal(|ui| {
widgets::hud(ui, "Connect account", theme::FAINT);
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
if widgets::ghost(ui, Icon::X, None, "Close", false).clicked() {
action = Some(SetupAction::Close);
}
});
});
ui.add_space(6.0);
ui.label(widgets::heading(&format!("Set up {name}"), 26.0));
ui.add_space(4.0);
ui.label(egui::RichText::new(format!("Add your {name} account to ArtCraft!")).size(14.0).family(theme::medium()));
ui.add_space(10.0);
let description = match provider {
Provider::Fal => format!("Paste your {name} API key to generate with your own {name} account. You can add all of your AI accounts to ArtCraft to use them all in one place and build the ultimate AI art tool."),
_ => format!("You can add your {name} account to ArtCraft by simply logging in. Use your credits and account directly within ArtCraft. You can add all of your AI accounts to ArtCraft to use them all in one place and build the ultimate AI art tool."),
};
ui.label(egui::RichText::new(description).size(13.0).color(theme::MUTED));
ui.add_space(16.0);
if provider == Provider::Fal {
text_field(ui, fal_key, true, "FAL API key");
ui.add_space(10.0);
}
let ready = provider != Provider::Fal || !fal_key.trim().is_empty();
if ui.add_enabled_ui(ready, |ui| widgets::button(ui, Some(Icon::ChevronRight), &format!("Set up {name}"), widgets::Kind::Primary, ui.available_width(), 40.0)).inner.clicked() {
action = Some(if provider == Provider::Fal { SetupAction::SaveFalKey } else { SetupAction::ConnectMidjourney });
}
});
if close && action.is_none() {
action = Some(SetupAction::Close);
}
action
}
/// What the accounts rows in Settings show.
pub struct AccountsView {
/// The Midjourney account's email ("" when unknown), when one is connected.
pub midjourney: Option<String>,
/// The saved fal key, masked.
pub fal_key_hint: Option<String>,
}
#[derive(Debug, PartialEq)]
pub enum AccountAction {
ConnectMidjourney,
DisconnectMidjourney,
SaveFalKey,
RemoveFalKey,
}
/// The provider accounts in Settings (`MidjourneyAccountBlock`, plus the fal key).
pub fn accounts_section(ui: &mut Ui, view: &AccountsView, fal_key: &mut String) -> Option<AccountAction> {
let mut action = None;
let detail = match view.midjourney.as_deref() {
Some("") => "Connected",
Some(email) => email,
None => "Not logged in",
};
let (label, kind) = if view.midjourney.is_some() { ("Disconnect", widgets::Kind::Danger) } else { ("Connect", widgets::Kind::Secondary) };
if account_row(ui, "Midjourney account", detail, |ui| widgets::button(ui, None, label, kind, 0.0, theme::CONTROL_H).clicked()) {
action = Some(if view.midjourney.is_some() { AccountAction::DisconnectMidjourney } else { AccountAction::ConnectMidjourney });
}
ui.add_space(10.0);
match &view.fal_key_hint {
Some(hint) => {
if account_row(ui, "FAL API key", hint, |ui| widgets::button(ui, None, "Remove", widgets::Kind::Danger, 0.0, theme::CONTROL_H).clicked()) {
action = Some(AccountAction::RemoveFalKey);
}
},
None => {
account_row(ui, "FAL API key", "Not set", |_| false);
ui.add_space(6.0);
ui.horizontal(|ui| {
let save_w = 90.0;
ui.allocate_ui(vec2(ui.available_width() - save_w - 8.0, theme::CONTROL_H), |ui| text_field(ui, fal_key, true, "Paste your FAL API key"));
let enabled = !fal_key.trim().is_empty();
if ui.add_enabled_ui(enabled, |ui| widgets::button(ui, None, "Save", widgets::Kind::Secondary, save_w, theme::CONTROL_H)).inner.clicked() {
action = Some(AccountAction::SaveFalKey);
}
});
},
}
action
}
/// A mono label over a detail line, with its button at the right. Returns the button's click.
fn account_row(ui: &mut Ui, label: &str, detail: &str, button: impl FnOnce(&mut Ui) -> bool) -> bool {
ui.horizontal(|ui| {
ui.vertical(|ui| {
ui.set_width(300.0);
widgets::hud(ui, label, theme::MUTED);
ui.label(egui::RichText::new(detail).size(13.0));
});
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), button).inner
})
.inner
}
/// The Midjourney sign-in window: a webview over a modal, checked every two seconds for the
/// login cookies (both auth tokens), which it hands back once present.
pub struct MidjourneyLogin {
context: WebContext,
webview: Option<WebView>,
error: Option<String>,
next_check: f64,
}
/// What came of the sign-in window.
pub enum LoginEvent {
Captured(CookieStore),
Closed,
}
impl MidjourneyLogin {
/// The window keeps its own browser profile under the data folder.
pub fn new(data_root: &Path) -> Self {
Self { context: WebContext::new(Some(data_root.join("webview").join("midjourney"))), webview: None, error: None, next_check: 0.0 }
}
pub fn show(&mut self, ctx: &egui::Context, window: &impl HasWindowHandle) -> Option<LoginEvent> {
let mut closed = false;
let mut slot = Rect::NOTHING;
let height = (ctx.content_rect().height() - 120.0).clamp(420.0, 800.0);
let dismissed = widgets::modal(ctx, Id::new("midjourney-login"), LOGIN_WIDTH, |ui| {
ui.horizontal(|ui| {
ui.label(widgets::heading("Login to Midjourney", 22.0));
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
if widgets::ghost(ui, Icon::X, None, "Close", false).clicked() {
closed = true;
}
});
});
ui.add_space(10.0);
let (rect, _) = ui.allocate_exact_size(vec2(ui.available_width(), height), Sense::hover());
ui.painter().rect_filled(rect, 0.0, theme::SUNKEN);
match &self.error {
Some(err) => {
let mut note = ui.new_child(egui::UiBuilder::new().max_rect(rect.shrink(24.0)).layout(egui::Layout::top_down(egui::Align::Center)));
note.add_space(rect.height() / 2.0 - 48.0);
note.label(egui::RichText::new(format!("Couldn't open the sign-in window: {err}")).size(13.0).color(theme::BAD));
},
None => icons::paint_spinner(ui, Rect::from_center_size(rect.center(), vec2(28.0, 28.0)), theme::MUTED),
}
slot = rect;
});
if closed || dismissed {
self.close();
return Some(LoginEvent::Closed);
}
if self.error.is_some() {
return None;
}
let bounds = physical_bounds(slot, ctx.pixels_per_point());
match &self.webview {
Some(webview) => {
if let Err(err) = webview.set_bounds(bounds) {
warn!("Couldn't move the sign-in window: {err}");
}
},
None => match self.open(window, bounds) {
Ok(webview) => self.webview = Some(webview),
Err(err) => self.error = Some(err.to_string()),
},
}
let now = ctx.input(|i| i.time);
ctx.request_repaint_after(std::time::Duration::from_secs_f64(COOKIE_CHECK_SECS));
if now < self.next_check {
return None;
}
self.next_check = now + COOKIE_CHECK_SECS;
let cookies = self.webview.as_ref().and_then(login_cookies)?;
self.close();
Some(LoginEvent::Captured(cookies))
}
/// Opens on a blank page, clears what an earlier sign-in left behind, then goes to Midjourney
/// (with the user agent the saved cookies will be replayed with).
fn open(&mut self, window: &impl HasWindowHandle, bounds: wry::Rect) -> wry::Result<WebView> {
let webview = WebViewBuilder::new_with_web_context(&mut self.context).with_user_agent(MIDJOURNEY_USER_AGENT).with_url("about:blank").with_bounds(bounds).build_as_child(window)?;
webview.clear_all_browsing_data()?;
webview.load_url(MIDJOURNEY_URL)?;
webview.focus()?;
Ok(webview)
}
fn close(&mut self) {
if let Some(webview) = self.webview.take() {
let _ = webview.focus_parent();
}
}
}
/// The login cookies, once the page is on Midjourney and both auth tokens are set.
fn login_cookies(webview: &WebView) -> Option<CookieStore> {
let host = webview.url().ok().and_then(|u| Url::parse(&u).ok()).and_then(|u| u.host_str().map(str::to_owned))?;
log::debug!("The Midjourney sign-in window is on {host}");
if host != "www.midjourney.com" && host != "midjourney.com" {
return None;
}
let mut store = CookieStore::empty();
for origin in MIDJOURNEY_ORIGINS {
let Ok(origin) = Url::parse(origin) else {
continue;
};
match webview.cookies_for_url(origin.as_str()) {
// The cookie's own text keeps its path, domain, Secure, HttpOnly, SameSite and expiry.
Ok(cookies) => cookies.iter().for_each(|c| {
store.apply_set_cookie_header(&c.to_string(), &origin);
}),
Err(err) => warn!("Couldn't read the sign-in cookies: {err}"),
}
}
cookie_store_has_auth_cookies(&store).then_some(store)
}
fn physical_bounds(rect: Rect, pixels_per_point: f32) -> wry::Rect {
let r = rect * pixels_per_point;
wry::Rect { position: wry::dpi::PhysicalPosition::new(r.left().round() as i32, r.top().round() as i32).into(), size: wry::dpi::PhysicalSize::new(r.width().round().max(1.0) as u32, r.height().round().max(1.0) as u32).into() }
}
+6 -1
View File
@@ -367,6 +367,11 @@ pub struct CostState {
impl CostState {
/// Requests an estimate for `fields` once they've been stable for a moment.
pub fn update(&mut self, ui: &Ui, backend: &Backend, modality: Modality, fields: &Map<String, Value>) {
self.update_with(ui, fields, |key, fields| backend.estimate_cost(modality, key, fields));
}
/// [`update`](Self::update), asking `request(key, fields)` for the estimate.
pub fn update_with(&mut self, ui: &Ui, fields: &Map<String, Value>, request: impl FnOnce(String, Map<String, Value>)) {
let key = Value::Object(fields.clone()).to_string();
if key == self.key {
return;
@@ -378,7 +383,7 @@ impl CostState {
self.key = key.clone();
self.credits = None;
self.pending_key = None;
backend.estimate_cost(modality, key, fields.clone());
request(key, fields.clone());
} else {
ui.ctx().request_repaint_after(std::time::Duration::from_secs_f64(COST_DEBOUNCE_SECS));
}
+3 -2
View File
@@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use crate::backend::media_cache::Lookup;
use crate::backend::providers::Provider;
use crate::backend::{EnqueueMeta, Modality};
use crate::feed::store::FeedStore;
use crate::feed::types::MediaKind;
@@ -125,7 +126,7 @@ impl CreateAudioPage {
let mut row_actions = Vec::new();
let mut slots = |ui: &mut Ui, slot: Slot| match (slot, model) {
(Slot::ToolbarLeft, _) => {
if let Some(id) = model_selector(ui, Id::new("audio-model"), models, model, false) {
if let Some((id, _)) = model_selector(ui, Id::new("audio-model"), models, model, false, None) {
picked_model = Some(id);
}
if let Some(m) = model {
@@ -318,7 +319,7 @@ impl CreateAudioPage {
return;
}
self.generating = true;
let meta = EnqueueMeta { prompt: self.settings.prompt.trim().to_owned(), model_id: m.id.clone(), batch_count: 1, ref_image: None };
let meta = EnqueueMeta { prompt: self.settings.prompt.trim().to_owned(), model_id: m.id.clone(), batch_count: 1, ref_image: None, provider: Provider::Artcraft };
env.backend.generate(Modality::Audio, self.fields(m, false), meta);
}
+54 -18
View File
@@ -1,10 +1,14 @@
//! Create Image (`TextToImage` + `PromptBoxImage`): prompt, reference images, model, aspect
//! ratio, resolution, quality and count; the cost; and the image feed.
//! Create Image (`TextToImage` + `PromptBoxImage`): prompt, reference images, model (and the
//! provider it runs on), aspect ratio, resolution, quality and count; the cost; and the image
//! feed.
use std::collections::HashMap;
use egui::{Id, Ui};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use crate::backend::providers::Provider;
use crate::backend::wire::Prompt;
use crate::backend::{EnqueueMeta, Modality};
use crate::feed::store::FeedStore;
@@ -29,6 +33,8 @@ pub struct ImageSettings {
pub resolution: Option<String>,
pub quality: Option<String>,
pub count: u16,
/// The provider picked for each model (the rest run on their default).
pub providers: HashMap<String, Provider>,
}
#[derive(Default)]
@@ -51,20 +57,32 @@ impl CreateImagePage {
pub fn ui(&mut self, ui: &mut Ui, env: &mut Env<'_>) {
let models = &env.catalog.image_page;
let model = common::selected_model(models, self.settings.model.as_deref(), DEFAULT_IMAGE_MODEL).cloned();
if let Some(m) = &model {
self.reconcile(m);
let fields = self.fields(m, true);
self.cost.update(ui, env.backend, Modality::Image, &fields);
let base = common::selected_model(models, self.settings.model.as_deref(), DEFAULT_IMAGE_MODEL).cloned();
let provider = base.as_ref().map_or(Provider::Artcraft, |m| self.provider(m));
// What the model can do on its provider (direct Midjourney: four images, no references).
let model = base.as_ref().map(|m| provider.adapt(m));
if let (Some(base), Some(m)) = (&base, &model) {
self.reconcile(base);
let mut fields = self.fields(m, true);
if provider != Provider::Artcraft {
fields.insert("provider".into(), json!(provider));
}
let backend = env.backend;
self.cost.update_with(ui, &fields, |key, fields| backend.estimate_provider_cost(provider, key, fields));
}
let mut feed = std::mem::take(&mut self.feed);
let mut box_height = self.box_height;
common::create_shell(ui, env, MediaKind::Image, "Create Image", "Add a prompt, then generate", &mut feed, &mut box_height, |ui, env| self.prompt_box(ui, env, model.as_ref()));
common::create_shell(ui, env, MediaKind::Image, "Create Image", "Add a prompt, then generate", &mut feed, &mut box_height, |ui, env| self.prompt_box(ui, env, model.as_ref(), provider));
self.feed = feed;
self.box_height = box_height;
}
fn prompt_box(&mut self, ui: &mut Ui, env: &mut Env<'_>, model: Option<&ModelInfo>) {
/// The provider `model` runs on: the user's pick, else its default.
fn provider(&self, model: &ModelInfo) -> Provider {
Provider::resolve(&model.id, self.settings.providers.get(&model.id).copied())
}
fn prompt_box(&mut self, ui: &mut Ui, env: &mut Env<'_>, model: Option<&ModelInfo>, provider: Provider) {
let max_images = model.map_or(0, |m| m.image_refs_max);
let deck = if max_images > 0 { DeckMode::References(DeckLimits { max_images, max_videos: 0, max_video_secs: None, max_audios: 0, max_audio_secs: None, library: env.signed_in }) } else { DeckMode::None };
let accepts: &[RefKind] = if max_images > 0 { &[RefKind::Image] } else { &[] };
@@ -77,8 +95,8 @@ impl CreateImagePage {
let mut picked_count = None;
let mut slots = |ui: &mut Ui, slot: Slot| match (slot, model) {
(Slot::ToolbarLeft, _) => {
if let Some(id) = model_selector(ui, Id::new("image-model"), models, model, false) {
picked_model = Some(id);
if let Some(pick) = model_selector(ui, Id::new("image-model"), models, model, false, Some(&settings.providers)) {
picked_model = Some(pick);
}
let Some(m) = model else {
return;
@@ -104,20 +122,23 @@ impl CreateImagePage {
_ => {},
};
let actions = prompt_box::show(ui, &mut settings.prompt, &self.refs, &mut self.box_state, &props, env.cache, &mut slots);
if let Some(id) = picked_model {
if let Some((id, picked_provider)) = picked_model {
if let Some(p) = picked_provider {
self.settings.providers.insert(id.clone(), p);
}
self.settings.model = Some(id);
}
if let Some(n) = picked_count {
self.settings.count = n;
}
for action in actions {
self.handle(action, env, model, max_images);
self.handle(action, env, model, max_images, provider);
}
}
fn handle(&mut self, action: PromptBoxAction, env: &mut Env<'_>, model: Option<&ModelInfo>, max_images: usize) {
fn handle(&mut self, action: PromptBoxAction, env: &mut Env<'_>, model: Option<&ModelInfo>, max_images: usize, provider: Provider) {
match action {
PromptBoxAction::Generate => self.generate(env, model),
PromptBoxAction::Generate => self.generate(env, model, provider),
PromptBoxAction::ClearAll => {
self.settings.prompt.clear();
self.refs.clear();
@@ -186,7 +207,7 @@ impl CreateImagePage {
}
}
fn generate(&mut self, env: &mut Env<'_>, model: Option<&ModelInfo>) {
fn generate(&mut self, env: &mut Env<'_>, model: Option<&ModelInfo>, provider: Provider) {
if self.generating || self.settings.prompt.trim().is_empty() {
return;
}
@@ -206,9 +227,24 @@ impl CreateImagePage {
env.toasts.error("Please wait for your references to finish uploading");
return;
}
// Other providers take the references as links (when the model takes references there).
let image_urls = if provider == Provider::Artcraft || model.image_refs_max == 0 {
Vec::new()
} else {
match self.refs.images.iter().map(|r| r.full_url.clone()).collect::<Option<Vec<String>>>() {
Some(urls) => urls,
None => {
env.toasts.error("A reference image is missing its link. Please add it again.");
return;
},
}
};
self.generating = true;
let meta = EnqueueMeta { prompt: self.settings.prompt.clone(), model_id: model.id.clone(), batch_count: u32::from(self.settings.count.max(1)), ref_image: self.refs.images.first().and_then(|r| r.preview.clone()) };
env.backend.generate(Modality::Image, self.fields(model, false), meta);
let meta = EnqueueMeta { prompt: self.settings.prompt.clone(), model_id: model.id.clone(), batch_count: u32::from(model.valid_batch(self.settings.count)), ref_image: self.refs.images.first().and_then(|r| r.preview.clone()), provider };
match provider {
Provider::Artcraft => env.backend.generate(Modality::Image, self.fields(model, false), meta),
_ => env.backend.generate_with_provider(provider, self.fields(model, false), image_urls, meta),
}
}
/// The OmniGen request (`GenerateImageRequest` after `api_fields`). `estimate` swaps
+3 -2
View File
@@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use crate::backend::wire::Prompt;
use crate::backend::providers::Provider;
use crate::backend::{EnqueueMeta, Modality};
use crate::feed::store::FeedStore;
use crate::feed::types::MediaKind;
@@ -126,7 +127,7 @@ impl CreateVideoPage {
let mut open_characters = false;
let mut slots = |ui: &mut Ui, slot: Slot| match (slot, model) {
(Slot::ToolbarLeft, _) => {
if let Some(id) = model_selector(ui, Id::new("video-model"), models, model, true) {
if let Some((id, _)) = model_selector(ui, Id::new("video-model"), models, model, true, None) {
picked_model = Some(id);
}
if let Some(m) = model {
@@ -316,7 +317,7 @@ impl CreateVideoPage {
}
self.generating = true;
let ref_image = self.refs.first_frame.as_ref().or(self.refs.images.first()).and_then(|r| r.preview.clone());
let meta = EnqueueMeta { prompt: self.settings.prompt.clone(), model_id: model.id.clone(), batch_count: u32::from(model.valid_batch(self.settings.count)), ref_image };
let meta = EnqueueMeta { prompt: self.settings.prompt.clone(), model_id: model.id.clone(), batch_count: u32::from(model.valid_batch(self.settings.count)), ref_image, provider: Provider::Artcraft };
let mut fields = self.fields(model, false);
if supports_characters(model) {
let tokens = mentioned_characters(&self.settings.prompt, env.characters, &self.character_picks);
+188 -38
View File
@@ -1,19 +1,31 @@
//! The embedded model selector (`ClassyModelSelector variant="embedded"`): a pill with the
//! maker's logo and the model's name, opening "Select Model". Long lists group by family, each
//! family opening a flyout of its variants on hover.
//! family opening a flyout of its variants on hover. Models that run on more than one provider
//! show the provider's logo, and hovering them opens "Select Provider".
use egui::{Align2, Id, Order, Popup, Rect, Ui, pos2};
use std::collections::HashMap;
use egui::{Align2, Id, Order, Popup, Rect, Stroke, Ui, pos2, vec2};
use crate::backend::providers::Provider;
use crate::models::{self, ModelInfo};
use crate::prompt_box::pickers;
use crate::theme;
use crate::ui::widgets::{self, Leading, RowLead};
use crate::ui::{creator_icons, icons};
const PANEL_WIDTH: f32 = 280.0;
const FLYOUT_WIDTH: f32 = 300.0;
const PROVIDER_FLYOUT_WIDTH: f32 = 220.0;
/// Hovering a model this long opens its providers (`providerTooltipDelayMs`).
const PROVIDER_DELAY_SECS: f64 = 0.3;
/// Draws the selector; returns the id of a newly picked model.
pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Option<&ModelInfo>, show_badges: bool) -> Option<String> {
/// What was picked: a model, and the provider when one was chosen with it.
pub type Pick = (String, Option<Provider>);
/// Draws the selector; returns a newly picked model (and provider). `providers` is the page's
/// provider choice per model (`None` where providers can't be chosen).
pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Option<&ModelInfo>, show_badges: bool, providers: Option<&HashMap<String, Provider>>) -> Option<Pick> {
let popup_id = ui.make_persistent_id(id).with("popup");
let open = Popup::is_id_open(ui.ctx(), popup_id);
let label = selected.map_or("Select model", |m| m.name.as_str());
@@ -26,6 +38,7 @@ pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Optio
}
});
let resp = resp.inner.on_hover_text("Model");
let rows = Rows { id, selected, show_badges, providers };
let mut picked = None;
let hovered_family_id = id.with("family");
pickers::popover(&resp, popup_id, |ui| {
@@ -42,12 +55,12 @@ pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Optio
let subtitle = if contains_selected { selected.map(|s| s.name.as_str()) } else { None };
let is_open = hovered.as_deref() == Some(family.name.as_str());
if family.models.len() == 1 {
let row = widgets::menu_row(ui, RowLead::Logo(&tex), &first.name, non_empty(&first.description), contains_selected, false, false);
if row.hovered() {
let (row_hovered, pick) = rows.row(ui, first);
if row_hovered {
hovered = None;
}
if row.clicked() {
picked = Some(first.id.clone());
if pick.is_some() {
picked = pick;
}
continue;
}
@@ -56,8 +69,8 @@ pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Optio
hovered = Some(family.name.clone());
}
if hovered.as_deref() == Some(family.name.as_str()) {
if let Some(id) = flyout(ui, id, row.rect, &family.name, &family.models, selected, show_badges) {
picked = Some(id);
if let Some(pick) = rows.family_flyout(ui, row.rect, &family.name, &family.models) {
picked = Some(pick);
}
}
}
@@ -65,8 +78,8 @@ pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Optio
},
None => {
for m in models {
if model_row(ui, m, selected, show_badges) {
picked = Some(m.id.clone());
if let (_, Some(pick)) = rows.row(ui, m) {
picked = Some(pick);
}
}
},
@@ -74,21 +87,109 @@ pub fn model_selector(ui: &mut Ui, id: Id, models: &[ModelInfo], selected: Optio
});
if picked.is_some() {
Popup::close_id(ui.ctx(), popup_id);
ui.data_mut(|d| d.remove::<Option<String>>(hovered_family_id));
ui.data_mut(|d| {
d.remove::<Option<String>>(hovered_family_id);
d.remove::<Option<(String, f64)>>(id.with("provider-hover"));
});
}
picked.filter(|p| selected.is_none_or(|s| &s.id != p))
picked.filter(|(model, provider)| provider.is_some() || selected.is_none_or(|s| &s.id != model))
}
/// A family's variants, to the right of its row.
fn flyout(ui: &Ui, id: Id, row: Rect, family: &str, list: &[&ModelInfo], selected: Option<&ModelInfo>, show_badges: bool) -> Option<String> {
/// What every row needs to know.
struct Rows<'a> {
id: Id,
selected: Option<&'a ModelInfo>,
show_badges: bool,
providers: Option<&'a HashMap<String, Provider>>,
}
impl Rows<'_> {
/// A family's variants, to the right of its row.
fn family_flyout(&self, ui: &Ui, row: Rect, family: &str, list: &[&ModelInfo]) -> Option<Pick> {
let mut picked = None;
egui::Area::new(self.id.with(("flyout", family))).order(Order::Tooltip).pivot(Align2::LEFT_TOP).fixed_pos(pos2(row.right() + 10.0, row.top() - 32.0)).show(ui.ctx(), |ui| {
widgets::popover_frame().show(ui, |ui| {
ui.set_width(FLYOUT_WIDTH);
widgets::menu_header(ui, family);
for m in list {
if let (_, Some(pick)) = self.row(ui, m) {
picked = Some(pick);
}
}
});
});
picked
}
/// One model: logo tile, name, description (and capability badges for video), and the
/// provider chip and flyout for models with a choice. Returns whether the row is hovered, and
/// a pick.
fn row(&self, ui: &mut Ui, m: &ModelInfo) -> (bool, Option<Pick>) {
let tex = creator_icons::texture(ui.ctx(), &m.creator);
let is_selected = self.selected.is_some_and(|s| s.id == m.id);
let badges = if self.show_badges { capability_badges(m) } else { String::new() };
let subtitle = match (m.description.is_empty(), badges.is_empty()) {
(true, true) => None,
(false, true) => Some(m.description.clone()),
(true, false) => Some(badges),
(false, false) => Some(format!("{} \u{b7} {badges}", m.description)),
};
let allowed = Provider::for_model(&m.id);
let Some(picks) = self.providers.filter(|_| allowed.len() > 1) else {
let row = widgets::menu_row(ui, RowLead::Logo(&tex), &m.name, subtitle.as_deref(), is_selected, false, false);
if row.hovered() {
ui.data_mut(|d| d.remove::<Option<(String, f64)>>(self.id.with("provider-hover")));
}
return (row.hovered(), row.clicked().then(|| (m.id.clone(), None)));
};
let current = Provider::resolve(&m.id, picks.get(&m.id).copied());
let hover_id = self.id.with("provider-hover");
let now = ui.input(|i| i.time);
let hovered: Option<(String, f64)> = ui.data(|d| d.get_temp(hover_id)).flatten();
let open = hovered.as_ref().is_some_and(|(id, since)| *id == m.id && now - since >= PROVIDER_DELAY_SECS);
let row = widgets::menu_row(ui, RowLead::Logo(&tex), &m.name, subtitle.as_deref(), false, false, open);
provider_chip(ui, row.rect, current, is_selected);
if row.hovered() && hovered.as_ref().is_none_or(|(id, _)| *id != m.id) {
ui.data_mut(|d| d.insert_temp(hover_id, Some((m.id.clone(), now))));
}
if row.hovered() && !open {
ui.ctx().request_repaint_after(std::time::Duration::from_secs_f64(PROVIDER_DELAY_SECS));
}
let mut picked = row.clicked().then(|| (m.id.clone(), None));
if open {
if let Some(provider) = provider_flyout(ui, self.id.with(("providers", &m.id)), row.rect, allowed, current) {
picked = Some((m.id.clone(), Some(provider)));
}
}
(row.hovered(), picked)
}
}
/// The provider's logo on a chip at the row's right; on the brand blue for the selected model.
fn provider_chip(ui: &Ui, row: Rect, provider: Provider, selected: bool) {
let chip = Rect::from_center_size(pos2(row.right() - 22.0, row.center().y), vec2(26.0, 20.0));
let fill = if selected { theme::fade(theme::ACCENT, 0.6) } else { theme::WASH };
ui.painter().rect_filled(chip, theme::RADIUS, fill);
let tex = creator_icons::texture(ui.ctx(), provider.icon());
egui::Image::new(&tex).tint(theme::INK).paint_at(ui, Rect::from_center_size(chip.center(), vec2(12.0, 12.0)));
}
/// "Select Provider" (`ProviderTooltipContent`), to the right of a model's row.
fn provider_flyout(ui: &Ui, id: Id, row: Rect, allowed: &[Provider], current: Provider) -> Option<Provider> {
let mut picked = None;
egui::Area::new(id.with(("flyout", family))).order(Order::Tooltip).pivot(Align2::LEFT_TOP).fixed_pos(pos2(row.right() + 10.0, row.top() - 32.0)).show(ui.ctx(), |ui| {
egui::Area::new(id).order(Order::Tooltip).pivot(Align2::LEFT_TOP).fixed_pos(pos2(row.right() + 10.0, row.top() - 32.0)).show(ui.ctx(), |ui| {
widgets::popover_frame().show(ui, |ui| {
ui.set_width(FLYOUT_WIDTH);
widgets::menu_header(ui, family);
for m in list {
if model_row(ui, m, selected, show_badges) {
picked = Some(m.id.clone());
ui.set_width(PROVIDER_FLYOUT_WIDTH);
widgets::menu_header(ui, "Select Provider");
for &provider in allowed {
let tex = creator_icons::texture(ui.ctx(), provider.icon());
let selected = provider == current;
let resp = widgets::menu_row(ui, RowLead::SmallLogo(&tex), provider.label(), None, selected, false, false);
if selected {
ui.painter().vline(resp.rect.left() + 1.5, resp.rect.y_range(), Stroke::new(3.0, theme::ACCENT));
}
if resp.clicked() {
picked = Some(provider);
}
}
});
@@ -96,20 +197,6 @@ fn flyout(ui: &Ui, id: Id, row: Rect, family: &str, list: &[&ModelInfo], selecte
picked
}
/// One model: logo tile, name, description (and capability badges for video).
fn model_row(ui: &mut Ui, m: &ModelInfo, selected: Option<&ModelInfo>, show_badges: bool) -> bool {
let tex = creator_icons::texture(ui.ctx(), &m.creator);
let is_selected = selected.is_some_and(|s| s.id == m.id);
let badges = if show_badges { capability_badges(m) } else { String::new() };
let subtitle = match (m.description.is_empty(), badges.is_empty()) {
(true, true) => None,
(false, true) => Some(m.description.clone()),
(true, false) => Some(badges),
(false, false) => Some(format!("{} \u{b7} {badges}", m.description)),
};
widgets::menu_row(ui, RowLead::Logo(&tex), &m.name, subtitle.as_deref(), is_selected, false, false).clicked()
}
/// The video picker's capability badges ("Audio Support", "Start/End", "Reference").
fn capability_badges(m: &ModelInfo) -> String {
let mut badges = Vec::new();
@@ -125,6 +212,69 @@ fn capability_badges(m: &ModelInfo) -> String {
badges.join(" \u{b7} ")
}
fn non_empty(s: &str) -> Option<&str> {
(!s.is_empty()).then_some(s)
/// The provider submenu headlessly (egui_kittest).
#[cfg(test)]
mod tests {
use egui_kittest::Harness;
use egui_kittest::kittest::Queryable;
use super::*;
struct State {
models: Vec<ModelInfo>,
picks: HashMap<String, Provider>,
picked: Vec<Pick>,
fonts_ready: bool,
}
#[test]
fn hovering_a_model_offers_its_providers() {
let mut harness = build();
harness.get_by_label("Seedream 4").click();
harness.run();
let row = harness.get_all_by_label("Nano Banana Pro").last().expect("the model's row");
row.hover();
for _ in 0..4 {
harness.step();
}
harness.run();
assert!(harness.query_by_label("ArtCraft").is_some(), "the submenu lists the providers");
harness.get_by_label("FAL").click();
harness.run();
assert_eq!(harness.state().picked, vec![("nano_banana_pro".to_owned(), Some(Provider::Fal))]);
}
#[test]
fn models_without_a_choice_have_no_submenu() {
let mut harness = build();
harness.get_by_label("Seedream 4").click();
harness.run();
harness.get_all_by_label("Seedream 4").last().expect("row").hover();
for _ in 0..4 {
harness.step();
}
harness.run();
assert!(harness.query_by_label("FAL").is_none(), "no provider list");
}
fn build() -> Harness<'static, State> {
let models = ["seedream_4", "nano_banana_pro"].iter().map(|id| ModelInfo { id: (*id).into(), name: models::display_name(id), creator: models::creator_for(id).into(), ..Default::default() }).collect();
let state = State { models, picks: HashMap::new(), picked: Vec::new(), fonts_ready: false };
let mut harness = Harness::builder().with_size(egui::vec2(1000.0, 700.0)).build_ui_state(
|ui, s: &mut State| {
if !s.fonts_ready {
theme::apply(ui.ctx());
s.fonts_ready = true;
return;
}
ui.add_space(500.0);
if let Some(pick) = model_selector(ui, Id::new("test-model"), &s.models, s.models.first(), false, Some(&s.picks)) {
s.picked.push(pick);
}
},
state,
);
harness.run();
harness
}
}
+6
View File
@@ -252,6 +252,8 @@ pub enum RowLead<'a> {
Aspect(Option<(f32, f32)>),
/// A 36 px tile with a model maker's logo (white texture, tinted).
Logo(&'a egui::TextureHandle),
/// A logo at icon size (the provider list).
SmallLogo(&'a egui::TextureHandle),
}
/// One option in a popover: a lead, a title, an optional muted subtitle, and a check when
@@ -284,6 +286,10 @@ pub fn menu_row(ui: &mut Ui, lead: RowLead<'_>, title: &str, subtitle: Option<&s
egui::Image::new(texture).tint(ink).paint_at(ui, Rect::from_center_size(tile.center(), vec2(18.0, 18.0)));
x += 46.0;
},
RowLead::SmallLogo(texture) => {
egui::Image::new(texture).tint(ink).paint_at(ui, Rect::from_center_size(pos2(x + 8.0, rect.center().y), vec2(15.0, 15.0)));
x += 26.0;
},
}
let title_font = FontId::new(13.5, theme::medium());
match subtitle {
@@ -6,3 +6,4 @@ pub mod image_generation_draft_context;
pub mod image_generation_draft_or_request;
pub mod image_generation_request;
pub mod providers;
pub mod router_image_request_to_artcraft_prompt;
@@ -1,21 +1,24 @@
//! The ArtCraft prompt record for an image request sent to another provider (Midjourney, fal),
//! created first so the results can be attributed to it.
use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
use artcraft_router::api::router_aspect_ratio::RouterAspectRatio;
use artcraft_router::api::router_image_model::RouterImageModel;
use artcraft_router::api::router_resolution::RouterResolution;
use artcraft_router::api::router_provider::RouterProvider;
use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder;
use enums::common::generation::common_aspect_ratio::CommonAspectRatio as EnumsAspectRatio;
use enums::common::generation::common_generation_mode::CommonGenerationMode;
use enums::common::generation::common_model_type::CommonModelType;
use enums::common::generation::common_resolution::CommonResolution as EnumsResolution;
use enums::common::generation_provider::GenerationProvider;
use uuid_utils::uuid::generate_random_uuid;
use crate::api::router_aspect_ratio::RouterAspectRatio;
use crate::api::router_image_model::RouterImageModel;
use crate::api::router_provider::RouterProvider;
use crate::api::router_resolution::RouterResolution;
use crate::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder;
pub fn router_image_request_to_artcraft_prompt(
request: &GenerateImageRequestBuilder,
) -> CreatePromptRequest {
CreatePromptRequest {
uuid_idempotency_token: generate_random_uuid(),
uuid_idempotency_token: uuid::Uuid::new_v4().to_string(),
positive_prompt: request.prompt.clone(),
negative_prompt: None,
model_type: image_model_to_common_model_type(request.model),
@@ -120,7 +123,7 @@ fn router_resolution_to_enums(res: RouterResolution) -> EnumsResolution {
#[cfg(test)]
mod tests {
use super::*;
use artcraft_router::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
use crate::client::request_mismatch_mitigation_strategy::RequestMismatchMitigationStrategy;
fn base_builder() -> GenerateImageRequestBuilder {
GenerateImageRequestBuilder {
@@ -165,7 +168,7 @@ mod tests {
#[test]
fn edit_mode_when_images_present() {
use artcraft_router::api::image_list_ref::ImageListRef;
use crate::api::image_list_ref::ImageListRef;
let builder = GenerateImageRequestBuilder {
image_inputs: Some(ImageListRef::Urls(vec!["https://example.com/img.jpg".to_string()])),
..base_builder()
@@ -14,14 +14,19 @@ browser_emulation.workspace = true
cloudflare_errors.workspace = true
cookie_store_wrapper = { workspace = true, features = ["wreq"] }
errors.workspace = true
uuid_utils.workspace = true
# External
anyhow.workspace = true
base64.workspace = true
chrono.workspace = true
futures.workspace = true
log.workspace = true
serde.workspace = true
serde_json.workspace = true
tempfile.workspace = true
tokio = { version = "1", features = ["full"] }
toml.workspace = true
url.workspace = true
wreq-util.workspace = true
wreq.workspace = true
@@ -0,0 +1,77 @@
//! Reading Midjourney's Firebase auth token (the `__Host-Midjourney.AuthUserTokenV3_i` cookie).
//! Its JWT payload is decoded WITHOUT verifying the signature: it's only used to learn the user
//! id, the email and when the token expires, never to trust anything.
use base64::Engine;
use base64::prelude::BASE64_URL_SAFE_NO_PAD;
use serde_json::Value;
use crate::credentials::cookie_names::AUTH_COOKIE_NAME_I;
use crate::credentials::midjourney_user_id::MidjourneyUserId;
/// The claim holding the Midjourney user id (the `singleplayer_{...}` channel id is built from
/// it).
const MIDJOURNEY_ID_CLAIM: &str = "midjourney_id";
/// Refresh this long before the token expires (seconds).
const REFRESH_MARGIN_SECS: i64 = 120;
/// The (unverified) payload of a `header.payload.signature` JWT.
pub fn jwt_claims(token: &str) -> Option<Value> {
let mut parts = token.split('.');
let (_header, payload, _signature) = (parts.next()?, parts.next()?, parts.next()?);
if parts.next().is_some() {
return None;
}
serde_json::from_slice(&BASE64_URL_SAFE_NO_PAD.decode(payload).ok()?).ok()
}
/// Whether the token is missing, unreadable, expired or about to expire.
pub 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 + REFRESH_MARGIN_SECS)
}
/// The Midjourney user id from the auth token's claims.
pub fn user_id_from_token(token: &str) -> Option<MidjourneyUserId> {
let claims = jwt_claims(token)?;
claims.get(MIDJOURNEY_ID_CLAIM).and_then(Value::as_str).filter(|id| !id.is_empty()).map(MidjourneyUserId::from_str)
}
/// The Midjourney user id from a raw `Cookie:` header (`"a=1; b=2"`), without hitting the
/// Cloudflare-gated index page.
pub fn extract_midjourney_user_id_from_cookie_header(cookie_header: &str) -> Option<MidjourneyUserId> {
let token = cookie_header.split(';').filter_map(|pair| pair.trim().split_once('=')).find(|(name, _)| *name == AUTH_COOKIE_NAME_I).map(|(_, value)| value)?;
user_id_from_token(token)
}
#[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_COOKIE_NAME_I, 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());
}
#[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));
}
fn token(exp: i64) -> String {
format!("header.{}.signature", BASE64_URL_SAFE_NO_PAD.encode(serde_json::json!({ "exp": exp }).to_string()))
}
}
@@ -1,13 +1,17 @@
//! The saved Midjourney login (`midjourney_cookies.toml`): the browser cookies captured at
//! sign-in, the account they belong to, and when they were last refreshed or failed.
use std::io::Write;
use std::path::Path;
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;
use crate::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.
@@ -0,0 +1,287 @@
//! The Midjourney login shared by the desktop apps: kept in one file, refreshed before it
//! expires, and turned into a [`MidjourneySession`] for each submit or poll.
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use anyhow::{anyhow, bail};
use browser_emulation::browser_profile::BrowserProfile;
use chrono::Utc;
use cookie_store_wrapper::cookie_store::CookieStore;
use serde_json::Value;
use tokio::sync::Mutex;
use url::Url;
use crate::credentials::auth_token::{extract_midjourney_user_id_from_cookie_header, jwt_claims, token_needs_refresh};
use crate::credentials::cookie_names::{AUTH_COOKIE_NAME_I, AUTH_COOKIE_NAME_R};
use crate::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
use crate::credentials::credential_file::{MidjourneyAccountInfo, MidjourneyCredential};
use crate::credentials::midjourney_user_id::MidjourneyUserId;
use crate::endpoints::refresh_token::{DEFAULT_FIREBASE_API_KEY, RefreshTokenArgs, RefreshTokenRequest, refresh_token};
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_error::MidjourneyError;
use crate::recipes::get_user_info::{GetUserInfoArgs, get_user_info};
/// Where the auth cookies are scoped.
pub const MIDJOURNEY_ORIGIN: &str = "https://www.midjourney.com/";
/// Token refresh and account lookups give up after this long.
const NETWORK_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Clone)]
pub struct MidjourneyCredentialManager {
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<()>>,
credential_path: PathBuf,
/// Older formats, removed on logout so they can't be migrated back.
legacy_paths: Vec<PathBuf>,
}
/// What a submit or poll needs: the cookies to send, whose account it is, and the browser to
/// present (the one the cookies were captured with).
pub struct MidjourneySession {
pub credential_id: String,
pub cookie_header: String,
pub user_id: MidjourneyUserId,
pub browser: BrowserProfile,
}
impl MidjourneyCredentialManager {
/// A manager that saves to `credential_path`, starting from `credential`.
pub fn new(credential_path: PathBuf, legacy_paths: Vec<PathBuf>, credential: Option<MidjourneyCredential>) -> Self {
Self { credential: Arc::new(RwLock::new(credential)), mutation: Arc::new(Mutex::new(())), credential_path, legacy_paths }
}
/// A manager for the login saved at `credential_path`, if there is one.
pub fn load(credential_path: PathBuf, legacy_paths: Vec<PathBuf>) -> anyhow::Result<Self> {
let credential = if credential_path.exists() {
let credential = MidjourneyCredential::load(&credential_path)?;
credential.save(&credential_path)?; // Persist a stable ID for hand-written files.
Some(credential)
} else {
None
};
Ok(Self::new(credential_path, legacy_paths, credential))
}
pub fn credential_path(&self) -> &PathBuf {
&self.credential_path
}
pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
Ok(self.snapshot()?.map(|credential| credential.cookie.cookies))
}
pub fn maybe_account_info(&self) -> anyhow::Result<Option<MidjourneyAccountInfo>> {
Ok(self.snapshot()?.and_then(|credential| credential.maybe_user_info))
}
pub fn session_appears_active(&self) -> anyhow::Result<bool> {
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_COOKIE_NAME_I), Utc::now().timestamp()))))
}
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 every format, so a later startup cannot migrate a logged-out account.
for path in self.legacy_paths.iter().chain([&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(())
}
/// 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);
};
if credential.cookie.maybe_relogin_required_since.is_some() {
return Ok(None);
}
if token_needs_refresh(credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_I), Utc::now().timestamp()) {
let Some(refresh) = credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_R).map(str::to_owned) else {
self.mark_relogin(&mut credential)?;
return Ok(None);
};
let refreshed = tokio::time::timeout(NETWORK_TIMEOUT, 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_COOKIE_NAME_I, response.id_token), (AUTH_COOKIE_NAME_R, 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(NETWORK_TIMEOUT, 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 }))
}
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())
}
}
/// Fills the account's id and email from the auth token's claims.
pub fn populate_account_info(credential: &mut MidjourneyCredential) {
let Some(claims) = credential.cookie.cookies.get_cookie_value(AUTH_COOKIE_NAME_I).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 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 base64::Engine;
use base64::prelude::BASE64_URL_SAFE_NO_PAD;
use tempfile::TempDir;
use super::*;
const TEST_USER_ID: &str = "00000000-0000-4000-8000-000000000001";
mod persistence {
use super::*;
#[tokio::test]
async fn expired_session_without_refresh_requires_login_and_new_login_clears_it() {
let (_dir, manager) = manager();
manager.save_login(cookies(10)).await.unwrap();
assert!(manager.session().await.unwrap().is_none());
assert!(!manager.session_appears_active().unwrap());
let dead = MidjourneyCredential::load(manager.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());
}
#[tokio::test]
async fn loads_a_saved_login_and_logout_removes_every_format() {
let (dir, manager) = manager();
manager.save_login(cookies(Utc::now().timestamp() + 3600)).await.unwrap();
let legacy = dir.path().join("midjourney_state.json");
std::fs::write(&legacy, "{}").unwrap();
let reloaded = MidjourneyCredentialManager::load(manager.credential_path().clone(), vec![legacy.clone()]).unwrap();
let session = reloaded.session().await.unwrap().unwrap();
assert_eq!(session.user_id.as_str(), TEST_USER_ID);
reloaded.clear_credentials().await.unwrap();
assert!(!legacy.exists());
assert!(!manager.credential_path().exists());
assert!(reloaded.session().await.unwrap().is_none());
}
#[test]
fn cookie_attributes_and_user_agent_survive_toml_round_trip() {
let (_dir, manager) = manager();
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 = manager.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 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 manager() -> (TempDir, MidjourneyCredentialManager) {
let dir = tempfile::tempdir().unwrap();
let manager = MidjourneyCredentialManager::new(dir.path().join("midjourney_cookies.toml"), Vec::new(), None);
(dir, manager)
}
fn cookies(exp: i64) -> CookieStore {
CookieStore::from_cookie_header(&format!("{AUTH_COOKIE_NAME_I}={}", 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()))
}
}
@@ -1,4 +1,7 @@
pub mod auth_token;
pub mod cookie_names;
pub mod cookie_store_has_auth_cookies;
pub mod credential_file;
pub mod credential_manager;
pub mod jwt_decode;
pub mod midjourney_user_id;
@@ -15,7 +15,7 @@ pub const MIDJOURNEY_USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X
/// 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.
/// with the login webview that captured the cookies.
pub fn midjourney_browser_profile() -> BrowserProfile {
BrowserProfile::safari_macos_with_user_agent(MIDJOURNEY_USER_AGENT)
}
@@ -1,2 +1,3 @@
pub mod get_image_url;
pub mod image_downloader_client;
pub mod image_downloader_client;
pub mod midjourney_browser_profile;
@@ -1,2 +1 @@
pub mod router_image_request_to_artcraft_prompt;
pub mod router_video_request_to_artcraft_prompt;
@@ -19,7 +19,7 @@ use crate::core::api_adapters::models::image::tauri_image_model_to_generation_mo
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::GenerateError;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::generate::common::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
use artcraft_router::generate::generate_image::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, convert_quality, convert_resolution};
use crate::core::commands::generate::generate_image::providers::artcraft_router::utils::map_media_files_to_urls::map_media_file_tokens_to_cdn_urls;
use crate::core::commands::generate::generate_image::tauri_generate_image_request::TauriGenerateImageRequest;
@@ -2,7 +2,7 @@ use crate::core::api_adapters::models::image::tauri_image_model_to_generation_mo
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::generate::common::router_image_request_to_artcraft_prompt::router_image_request_to_artcraft_prompt;
use artcraft_router::generate::generate_image::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;
@@ -1,337 +1,126 @@
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;
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;
//! The app's Midjourney login: the shared manager (`midjourney_client`, also used by the native
//! app) kept in the app's credentials directory, plus the one-time migration from the legacy
//! `midjourney_state.json`.
use std::ops::Deref;
use cookie_store_wrapper::cookie_store::CookieStore;
use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
use midjourney_client::credentials::credential_file::{MidjourneyAccountInfo, MidjourneyCredential};
use midjourney_client::credentials::credential_manager::{MIDJOURNEY_ORIGIN, MidjourneyCredentialManager as SharedCredentialManager, populate_account_info};
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/";
pub use midjourney_client::credentials::credential_manager::MidjourneySession;
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
use crate::services::midjourney::state::serializable_midjourney_state::SerializableMidjourneyState;
#[derive(Clone)]
pub struct MidjourneyCredentialManager {
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 MidjourneyCredentialManager(SharedCredentialManager);
pub struct MidjourneySession {
pub credential_id: String,
pub cookie_header: String,
pub user_id: MidjourneyUserId,
pub browser: BrowserProfile,
impl Deref for MidjourneyCredentialManager {
type Target = SharedCredentialManager;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl MidjourneyCredentialManager {
pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self {
Self { credential: Arc::new(RwLock::new(None)), mutation: Arc::new(Mutex::new(())), app_data_root: app_data_root.clone() }
let dir = app_data_root.credentials_dir();
Self(SharedCredentialManager::new(dir.get_midjourney_credential_path(), vec![dir.get_midjourney_state_path()], None))
}
pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self {
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(manager) => manager,
Err(err) => {
log::warn!("Could not load Midjourney credentials: {}", err);
Self::initialize_empty(app_data_root)
},
Err(err) => log::warn!("Could not load Midjourney credentials: {}", err),
}
manager
}
pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
Ok(self.snapshot()?.map(|credential| credential.cookie.cookies))
}
pub fn maybe_copy_user_info(&self) -> anyhow::Result<Option<MidjourneyUserInfo>> {
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 session_appears_active(&self) -> anyhow::Result<bool> {
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 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(())
}
/// 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);
};
if credential.cookie.maybe_relogin_required_since.is_some() {
return Ok(None);
}
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 }))
}
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()
Ok(self.maybe_account_info()?.map(|info| MidjourneyUserInfo { user_id: info.maybe_user_id.map(MidjourneyUserId::from_string), email: info.maybe_email, websocket_token: None }))
}
}
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));
fn load_or_migrate(root: &AppDataRoot) -> anyhow::Result<MidjourneyCredentialManager> {
let dir = root.credentials_dir();
let (path, legacy_path) = (dir.get_midjourney_credential_path(), dir.get_midjourney_state_path());
if path.exists() || !legacy_path.exists() {
return Ok(MidjourneyCredentialManager(SharedCredentialManager::load(path, vec![legacy_path])?));
}
let legacy_path = root.credentials_dir().get_midjourney_state_path();
if !legacy_path.exists() {
return Ok(None);
}
let legacy: SerializableMidjourneyState = serde_json::from_str(&std::fs::read_to_string(legacy_path)?)?;
let legacy: SerializableMidjourneyState = serde_json::from_str(&std::fs::read_to_string(&legacy_path)?)?;
let Some(cookies) = legacy.user_cookies else {
return Ok(None);
return Ok(MidjourneyCredentialManager::initialize_empty(root));
};
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))
Ok(MidjourneyCredentialManager(SharedCredentialManager::new(path, vec![legacy_path], Some(credential))))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine};
use chrono::Utc;
use tempfile::TempDir;
use super::*;
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
const TEST_USER_ID: &str = "00000000-0000-4000-8000-000000000001";
const AUTH_ID: &str = "__Host-Midjourney.AuthUserTokenV3_i";
const AUTH_REFRESH: &str = "__Host-Midjourney.AuthUserTokenV3_r";
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());
}
#[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());
}
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()));
}
#[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());
}
fn root() -> (TempDir, AppDataRoot) {
@@ -341,10 +130,6 @@ mod tests {
(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()))
@@ -1,4 +1,3 @@
pub (super) mod serializable_midjourney_state;
pub mod midjourney_credential_manager;
pub mod midjourney_user_info;
pub mod midjourney_credential;
@@ -169,7 +169,7 @@ 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 midjourney_client::credentials::credential_file::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;
@@ -1,62 +0,0 @@
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());
}
}
@@ -1,3 +1 @@
pub mod download_midjourney_image;
pub mod midjourney_browser_profile;
pub mod extract_midjourney_user_id_from_cookies;
@@ -1,4 +1,4 @@
use crate::services::midjourney::utils::midjourney_browser_profile::MIDJOURNEY_USER_AGENT;
use midjourney_client::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;