mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-26 03:39:28 -04:00
Compare commits
14
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
791f0ba21a | ||
|
|
06e82eb9dc | ||
|
|
ff443562db | ||
|
|
a95bc53eb5 | ||
|
|
6a3693a7ce | ||
|
|
e3c99add5d | ||
|
|
8c3f880655 | ||
|
|
08300cb47c | ||
|
|
0dc0a2f597 | ||
|
|
f91f1965a4 | ||
|
|
cc1ed0eb76 | ||
|
|
93e9d423cb | ||
|
|
e6267eeb23 | ||
|
|
4e28d05e51 |
@@ -26,14 +26,23 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: install Rust stable
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: install wreq dependencies
|
||||
run: |
|
||||
choco install cmake -y
|
||||
choco install strawberryperl -y
|
||||
choco install pkgconfiglite -y
|
||||
choco install llvm -y
|
||||
choco install nasm -y
|
||||
shell: cmd
|
||||
|
||||
- name: setup node
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: lts/*
|
||||
|
||||
- name: install Rust stable
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: install frontend dependencies
|
||||
run: |
|
||||
# Navigate to the frontend directory to install dependencies
|
||||
@@ -83,6 +92,8 @@ jobs:
|
||||
- name: build and publish
|
||||
uses: tauri-apps/tauri-action@v0
|
||||
env:
|
||||
# NB(bt): This is needed by wreq. Hopefully it doesn't kill Tauri
|
||||
RUSTFLAGS: "-C target-feature=+crt-static"
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Use the SQLx query cache for SQLite
|
||||
SQLX_OFFLINE: true
|
||||
|
||||
-12
@@ -1,12 +0,0 @@
|
||||
{
|
||||
"db_name": "MySQL",
|
||||
"query": "\nINSERT INTO media_files\nSET\n token = ?,\n\n media_class = ?,\n media_type = ?,\n\n origin_category = ?,\n origin_product_category = ?,\n\n maybe_engine_category = ?,\n maybe_animation_type = ?,\n\n maybe_mime_type = ?,\n file_size_bytes = ?,\n\n checksum_sha2 = ?,\n\n maybe_title = ?,\n maybe_duration_millis = ?,\n\n public_bucket_directory_hash = ?,\n maybe_public_bucket_prefix = ?,\n maybe_public_bucket_extension = ?,\n\n maybe_creator_user_token = ?,\n maybe_creator_anonymous_visitor_token = ?,\n creator_ip_address = ?,\n\n creator_set_visibility = ?,\n\n maybe_scene_source_media_file_token = ?,\n is_intermediate_system_file = ?,\n\n maybe_creator_file_synthetic_id = ?,\n maybe_creator_category_synthetic_id = ?,\n\n is_user_upload = TRUE,\n\n maybe_origin_model_type = NULL,\n maybe_origin_model_token = NULL\n ",
|
||||
"describe": {
|
||||
"columns": [],
|
||||
"parameters": {
|
||||
"Right": 23
|
||||
},
|
||||
"nullable": []
|
||||
},
|
||||
"hash": "28cac5883464613d29bedbf7308a71c258f1bf4abb8f5ca3ae320db34c137955"
|
||||
}
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"db_name": "MySQL",
|
||||
"query": "\nINSERT INTO media_files\nSET\n token = ?,\n\n media_class = ?,\n media_type = ?,\n\n origin_category = ?,\n origin_product_category = ?,\n\n maybe_engine_category = ?,\n maybe_animation_type = ?,\n\n maybe_mime_type = ?,\n file_size_bytes = ?,\n\n checksum_sha2 = ?,\n\n maybe_prompt_token = ?,\n\n maybe_title = ?,\n maybe_duration_millis = ?,\n\n public_bucket_directory_hash = ?,\n maybe_public_bucket_prefix = ?,\n maybe_public_bucket_extension = ?,\n\n maybe_creator_user_token = ?,\n maybe_creator_anonymous_visitor_token = ?,\n creator_ip_address = ?,\n\n creator_set_visibility = ?,\n\n maybe_scene_source_media_file_token = ?,\n is_intermediate_system_file = ?,\n\n maybe_creator_file_synthetic_id = ?,\n maybe_creator_category_synthetic_id = ?,\n\n is_user_upload = TRUE,\n\n maybe_origin_model_type = NULL,\n maybe_origin_model_token = NULL\n ",
|
||||
"describe": {
|
||||
"columns": [],
|
||||
"parameters": {
|
||||
"Right": 24
|
||||
},
|
||||
"nullable": []
|
||||
},
|
||||
"hash": "e86494c0a53247903aa58009e2a26c4db90a89e9f6ea805ffb3929faf69a431d"
|
||||
}
|
||||
Generated
+581
-110
File diff suppressed because it is too large
Load Diff
+7
-1
@@ -10,6 +10,7 @@ members = [
|
||||
"crates/api_clients/fal_client",
|
||||
"crates/api_clients/google_drive_common",
|
||||
"crates/api_clients/google_sign_in",
|
||||
"crates/api_clients/midjourney_client",
|
||||
"crates/api_clients/openai_sora_client",
|
||||
"crates/api_clients/openart_client",
|
||||
"crates/api_clients/runwayml_client",
|
||||
@@ -42,6 +43,7 @@ members = [
|
||||
"crates/lib/caching/redis_caching",
|
||||
"crates/lib/cloudflare_errors",
|
||||
"crates/lib/concurrency",
|
||||
"crates/lib/cookie_store",
|
||||
"crates/lib/cookies",
|
||||
"crates/lib/data/collections",
|
||||
"crates/lib/data/crockford",
|
||||
@@ -199,7 +201,8 @@ log = "0.4.27"
|
||||
once_cell = "1.21.3"
|
||||
rand = "0.9.1"
|
||||
regex = "1.11.1"
|
||||
reqwest = { version = "0.12.15", features = ["json", "cookies"] }
|
||||
reqwest = { version = "0.12.22", features = ["json", "cookies"] }
|
||||
reqwest-websocket = { version = "0.5.0" }
|
||||
serde = { version = "1.0.219", features = ["derive"] }
|
||||
serde_derive = "1.0.219"
|
||||
serde_json = "1.0.140"
|
||||
@@ -218,6 +221,7 @@ artcraft_api_defs = { path = "crates/schema/public/api/artcraft_api_defs" }
|
||||
build_metadata = { path = "crates/lib/build_metadata" }
|
||||
cloud_storage = { path = "crates/api_clients/cloud_storage" }
|
||||
cloudflare_errors = { path = "crates/lib/cloudflare_errors" }
|
||||
cookie_store = { path = "crates/lib/cookie_store" }
|
||||
crockford = { path = "crates/lib/data/crockford" }
|
||||
enums = { path = "crates/schema/public/enums" }
|
||||
enums_public = { path = "crates/schema/public/enums_public" }
|
||||
@@ -230,6 +234,7 @@ images = { path = "crates/lib/files/images" }
|
||||
jwt_light = { path = "crates/lib/jwt_light" }
|
||||
memory_caching = { path = "crates/lib/caching/memory_caching" }
|
||||
memory_store = { path = "crates/lib/caching/memory_store" }
|
||||
midjourney_client = { path = "crates/api_clients/midjourney_client" }
|
||||
mimetypes = { path = "crates/lib/files/mimetypes" }
|
||||
openai_sora_client = { path = "crates/api_clients/openai_sora_client" }
|
||||
primitives = { path = "crates/lib/data/primitives" }
|
||||
@@ -238,5 +243,6 @@ sqlite_tasks = { path = "crates/schema/database/sqlite_tasks" }
|
||||
storyteller_client = { path = "crates/api_clients/storyteller_client" }
|
||||
subprocess_common = { path = "crates/lib/subprocess_common" }
|
||||
test_data = { path = "crates/lib/test_data" }
|
||||
testing = { path = "crates/lib/testing" }
|
||||
tokens = { path = "crates/schema/public/tokens" }
|
||||
web_base64 = { path = "crates/lib/data/web_base64" }
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
-- noinspection SqlDialectInspectionForFile
|
||||
--
|
||||
-- NB: This version needs to be bumped even if just adding or changing comments!
|
||||
-- Failure to do so will cause Windows to segfault at start. (This is horrible!)
|
||||
--
|
||||
-- NB: `app_state_dir.rs` contains the file version number.
|
||||
--
|
||||
-- Migration history:
|
||||
-- tasks_v1.sqlite - initial version
|
||||
-- tasks_v2.sqlite - added model_type (nullable)
|
||||
-- tasks_v3.sqlite - added frontend_caller, comments
|
||||
|
||||
CREATE TABLE tasks (
|
||||
-- Task auto-incrementing primary key.
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
[package]
|
||||
name = "midjourney_client"
|
||||
edition = "2024"
|
||||
version = "0.0.1"
|
||||
publish = false
|
||||
|
||||
[package.metadata.cargo-udeps.ignore]
|
||||
normal = ["workspace-build-acceleration"]
|
||||
|
||||
[lib]
|
||||
name = "midjourney_client"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration.workspace = true
|
||||
|
||||
# Internal
|
||||
cloudflare_errors.workspace = true
|
||||
cookie_store.workspace = true
|
||||
errors.workspace = true
|
||||
|
||||
# External
|
||||
anyhow.workspace = true
|
||||
futures.workspace = true
|
||||
log.workspace = true
|
||||
#reqwest-websocket.workspace = true
|
||||
#reqwest.workspace = true
|
||||
#reqwest-impersonate = { version = "0", features = ["websocket"] }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
#tokio.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
# NB: wreq gets past fingerprinting, but we have to deal with platform-specific build issues.
|
||||
#
|
||||
# Mac:
|
||||
# brew install --formula cmake pkg-config llvm
|
||||
# Windows:
|
||||
# choco install cmake -y
|
||||
# choco install strawberryperl -y
|
||||
# choco install pkgconfiglite -y
|
||||
# choco install llvm -y
|
||||
# choco install nasm -y
|
||||
wreq = { version = "5", features = ["json"] }
|
||||
wreq-util = "2"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
# NB: We need to parse HTML to get out the websocket token
|
||||
dom_query = "0.20.2"
|
||||
|
||||
[dev-dependencies]
|
||||
filesys.workspace = true
|
||||
test_data.workspace = true
|
||||
testing.workspace = true
|
||||
@@ -0,0 +1,17 @@
|
||||
|
||||
/// Multiple third party Midjourney clients make this configurable.
|
||||
/// Not sure why, I'm just cargo culting it.
|
||||
#[derive(Clone)]
|
||||
pub enum MidjourneyHostname {
|
||||
Standard,
|
||||
Custom(String),
|
||||
}
|
||||
|
||||
impl MidjourneyHostname {
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
Self::Standard => "www.midjourney.com",
|
||||
Self::Custom(domain) => domain,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod midjourney_hostname;
|
||||
pub mod midjourney_websocket_client;
|
||||
@@ -0,0 +1,6 @@
|
||||
|
||||
/// One of the apparent authentication cookies used by Midjourney.
|
||||
pub (crate) const AUTH_COOKIE_NAME_R : &str = "__Host-Midjourney.AuthUserTokenV3_r";
|
||||
|
||||
/// One of the apparent authentication cookies used by Midjourney.
|
||||
pub (crate) const AUTH_COOKIE_NAME_I : &str = "__Host-Midjourney.AuthUserTokenV3_i";
|
||||
@@ -0,0 +1,8 @@
|
||||
use crate::credentials::cookie_names::{AUTH_COOKIE_NAME_I, AUTH_COOKIE_NAME_R};
|
||||
use cookie_store::cookie_store::CookieStore;
|
||||
|
||||
/// Returns true if the given cookie store has the necessary authentication cookies.
|
||||
pub fn cookie_store_has_auth_cookies(cookie_store: &CookieStore) -> bool {
|
||||
cookie_store.has_cookie(AUTH_COOKIE_NAME_I)
|
||||
&& cookie_store.has_cookie(AUTH_COOKIE_NAME_R)
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
|
||||
/*
|
||||
{
|
||||
"name": "u3869671173",
|
||||
"picture": "https://lh3.googleusercontent.com/a/ACg8ocKw0aaXKYSqy-Xkkn69ydIJFGnJA37S13XvjNY4d9xCzUzKUw=s96-c",
|
||||
"midjourney_id": "f8a57ac3-e416-4dd4-9be8-2c4223691b01",
|
||||
"iss": "https://securetoken.google.com/authjourney",
|
||||
"aud": "authjourney",
|
||||
"auth_time": 1754970849,
|
||||
"user_id": "bS9idl73DlSbE61KnCyKpDfazCX2",
|
||||
"sub": "bS9idl73DlSbE61KnCyKpDfazCX2",
|
||||
"iat": 1754970849,
|
||||
"exp": 1754974449,
|
||||
"email": "victorvoice89@gmail.com",
|
||||
"email_verified": true,
|
||||
"firebase": {
|
||||
"identities": {
|
||||
"google.com": [
|
||||
"116459300198619888836"
|
||||
],
|
||||
"email": [
|
||||
"victorvoice89@gmail.com"
|
||||
]
|
||||
},
|
||||
"sign_in_provider": "google.com"
|
||||
}
|
||||
}
|
||||
*/
|
||||
|
||||
/*
|
||||
*/
|
||||
@@ -0,0 +1,22 @@
|
||||
|
||||
/// The user ID used for several midjourney requests.
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
|
||||
pub struct MidjourneyUserId(pub String);
|
||||
|
||||
impl MidjourneyUserId {
|
||||
pub fn from_str(user_id: &str) -> Self {
|
||||
Self(user_id.to_string())
|
||||
}
|
||||
|
||||
pub fn from_string(user_id: String) -> Self {
|
||||
Self(user_id)
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub fn to_string(&self) -> String {
|
||||
self.0.to_string()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod cookie_names;
|
||||
pub mod cookie_store_has_auth_cookies;
|
||||
pub mod jwt_decode;
|
||||
pub mod midjourney_user_id;
|
||||
@@ -0,0 +1,95 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
|
||||
use log::error;
|
||||
use serde::Deserialize;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
/// This grabs the pre-rendered HTML page, which contains the websocket token and other details.
|
||||
pub struct GetIndexPageRequest {
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
pub async fn get_index_page_html(req: GetIndexPageRequest) -> Result<String, MidjourneyError> {
|
||||
let url = format!("https://{}/", req.hostname.as_str());
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox139)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let cookie_header = req.cookie_header.trim();
|
||||
|
||||
if cookie_header.len() < 20 {
|
||||
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
|
||||
return Err(MidjourneyClientError::CookieTooShort.into());
|
||||
}
|
||||
|
||||
// NB: missing headers that were in the browser request:
|
||||
// -H 'sec-ch-ua-platform: "macOS"' \
|
||||
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
|
||||
|
||||
let mut http_request = client.get(url)
|
||||
.header("cookie", cookie_header)
|
||||
.header("Referrer-Policy", "origin-when-cross-origin")
|
||||
.header("accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7")
|
||||
.header("accept-language", "en-US,en;q=0.8")
|
||||
.header("content-type", "application/json")
|
||||
.header("priority", "u=0, i")
|
||||
.header("sec-ch-ua-mobile", "?0")
|
||||
.header("sec-fetch-dest", "document")
|
||||
.header("sec-fetch-mode", "navigate")
|
||||
.header("sec-fetch-site", "none")
|
||||
.header("sec-fetch-user", "?1")
|
||||
.header("upgrade-insecure-requests", "1");
|
||||
|
||||
let http_request = http_request
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let response = client.execute(http_request)
|
||||
.await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
let response_body = response.text().await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
|
||||
return Err(MidjourneyApiError::CloudflareError(err).into());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(response_body)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageRequest};
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
|
||||
let result = get_index_page_html(GetIndexPageRequest {
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}\n\n", result);
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
|
||||
use log::error;
|
||||
use serde::Deserialize;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
const DEFAULT_PAGE_SIZE : u64 = 1000;
|
||||
|
||||
/// This returns Midjourney's job/media list
|
||||
pub struct ImagineRequest<'a> {
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
pub user_id: &'a MidjourneyUserId,
|
||||
pub page_size: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ImagineResponse {
|
||||
pub items: Vec<ImagineItem>,
|
||||
pub cursor: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ImagineItem {
|
||||
pub id: Option<String>,
|
||||
|
||||
/// This is the prompt, e.g.
|
||||
/// `cute puppy wearing a polar bear suit --ar 4:3 --raw --v 7.0`
|
||||
pub full_command: Option<String>,
|
||||
|
||||
pub job_type: Option<MidjourneyJobType>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum MidjourneyJobType {
|
||||
V6Diffusion,
|
||||
V6p1Diffusion,
|
||||
V6p1RawDiffusion,
|
||||
V7Diffusion,
|
||||
V7DraftDiffusion,
|
||||
V7DraftRawDiffusion,
|
||||
V7RawDiffusion,
|
||||
Vid11I2vRenderAJointVideo,
|
||||
Vid11I2vRenderBJointVideo,
|
||||
Vid11I2vStartEndAVideo,
|
||||
Vid11I2vStartEndBVideo,
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl MidjourneyJobType {
|
||||
pub fn from_str(s: &str) -> MidjourneyJobType {
|
||||
match s {
|
||||
"v6_diffusion" => MidjourneyJobType::V6Diffusion,
|
||||
"v6-1_diffusion" => MidjourneyJobType::V6p1Diffusion,
|
||||
"v6-1_raw_diffusion" => MidjourneyJobType::V6p1RawDiffusion,
|
||||
"v7_diffusion" => MidjourneyJobType::V7Diffusion,
|
||||
"v7_draft_diffusion" => MidjourneyJobType::V7DraftDiffusion,
|
||||
"v7_draft_raw_diffusion" => MidjourneyJobType::V7DraftRawDiffusion,
|
||||
"v7_raw_diffusion" => MidjourneyJobType::V7RawDiffusion,
|
||||
"vid_1.1_i2v_render_a_joint_video" => MidjourneyJobType::Vid11I2vRenderAJointVideo,
|
||||
"vid_1.1_i2v_render_b_joint_video" => MidjourneyJobType::Vid11I2vRenderBJointVideo,
|
||||
"vid_1.1_i2v_start_end_a_video" => MidjourneyJobType::Vid11I2vStartEndAVideo,
|
||||
"vid_1.1_i2v_start_end_b_video" => MidjourneyJobType::Vid11I2vStartEndBVideo,
|
||||
other => MidjourneyJobType::Other(other.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn imagine(req: ImagineRequest<'_>) -> Result<ImagineResponse, MidjourneyError> {
|
||||
let cookie_header = req.cookie_header.trim();
|
||||
|
||||
if cookie_header.len() < 20 {
|
||||
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
|
||||
return Err(MidjourneyClientError::CookieTooShort.into());
|
||||
}
|
||||
|
||||
let page_size = req.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
|
||||
|
||||
let referer = format!("https://{}/imagine", req.hostname.as_str());
|
||||
|
||||
let url = format!("https://{}/api/imagine?user_id={}&page_size={}",
|
||||
req.hostname.as_str(),
|
||||
req.user_id.as_str(),
|
||||
page_size,
|
||||
);
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox139)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
// NB: missing headers that were in the browser request:
|
||||
// -H 'sec-ch-ua-platform: "macOS"' \
|
||||
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
|
||||
|
||||
let mut http_request = client.get(url)
|
||||
.header("cookie", cookie_header)
|
||||
.header("Referer", &referer)
|
||||
.header("Referrer-Policy", "origin-when-cross-origin")
|
||||
.header("accept", "*/*")
|
||||
.header("accept-encoding", "gzip, deflate, br, zstd")
|
||||
.header("accept-language", "en-US,en;q=0.8")
|
||||
.header("content-type", "application/json")
|
||||
.header("priority", "u=1, i")
|
||||
.header("sec-ch-ua-mobile", "?0")
|
||||
.header("sec-fetch-dest", "empty")
|
||||
.header("sec-fetch-mode", "cors")
|
||||
.header("sec-fetch-site", "same-origin")
|
||||
.header("x-csrf-protection", "1");
|
||||
|
||||
let http_request = http_request
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let response = client.execute(http_request)
|
||||
.await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
let response_body = response.text().await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
|
||||
return Err(MidjourneyApiError::CloudflareError(err).into());
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
{
|
||||
"data": [
|
||||
{
|
||||
"id": "a2339f14-ff71-43e4-a69f-8da15e7ba871",
|
||||
"enqueue_time": "2025-08-14T17:21:28.403560+00:00",
|
||||
"parent_id": null,
|
||||
"video_segments": null,
|
||||
"rating": null,
|
||||
"job_type": "v7_diffusion",
|
||||
"event_type": "diffusion",
|
||||
"parent_grid": null,
|
||||
"full_command": "a modern n64 console",
|
||||
"batch_size": 4,
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"published": true,
|
||||
"shown": true,
|
||||
"user_hidden": false,
|
||||
"template": null,
|
||||
"vid_framecount": null
|
||||
},
|
||||
],
|
||||
"cursor": "encrypted cursor",
|
||||
"checkpoint": "encrypted something ???"
|
||||
}
|
||||
*/
|
||||
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawImagineResponse {
|
||||
data: Vec<RawImagineItem>,
|
||||
cursor: Option<String>,
|
||||
//pub checkpoint: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawImagineItem {
|
||||
// The job id
|
||||
id: Option<String>,
|
||||
|
||||
// The prompt plus `--` args.
|
||||
full_command: Option<String>,
|
||||
|
||||
// Number of images, typically 4
|
||||
batch_size: Option<u8>,
|
||||
|
||||
// Some possible values:
|
||||
// "v6_diffusion"
|
||||
// "v7_diffusion"
|
||||
// "v7_draft_raw_diffusion"
|
||||
// "v7_raw_diffusion"
|
||||
// "vid_1.1_i2v_render_b_joint_video",
|
||||
// "vid_1.1_i2v_start_end_a_video"
|
||||
// "vid_1.1_i2v_start_end_b_video"
|
||||
job_type: Option<String>,
|
||||
|
||||
// Some possible values:
|
||||
// "diffusion"
|
||||
// "draft"
|
||||
// "variation"
|
||||
// "video_diffusion"
|
||||
// "video_start_end"
|
||||
event_type: Option<String>,
|
||||
}
|
||||
|
||||
let response : RawImagineResponse = serde_json::from_str(&response_body)
|
||||
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
|
||||
|
||||
Ok(ImagineResponse {
|
||||
cursor: response.cursor.clone(),
|
||||
items: response.data
|
||||
.into_iter()
|
||||
.map(|item| {
|
||||
ImagineItem {
|
||||
id: item.id,
|
||||
full_command: item.full_command,
|
||||
job_type: item.job_type.as_ref()
|
||||
.map(|jt| MidjourneyJobType::from_str(jt))
|
||||
}
|
||||
}).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use crate::endpoints::imagine::{imagine, ImagineRequest};
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
let user_id = read_to_trimmed_string("/Users/bt/secrets/midjourney/user_id.txt")?;
|
||||
let user_id = MidjourneyUserId::from_string(user_id);
|
||||
|
||||
let result = imagine(ImagineRequest {
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
user_id: &user_id,
|
||||
page_size: None,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}\n\n", result);
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod get_index_page_html;
|
||||
pub mod imagine;
|
||||
pub mod storage_list;
|
||||
pub mod submit_job;
|
||||
pub mod user_connections;
|
||||
@@ -0,0 +1,147 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
|
||||
use log::error;
|
||||
use serde::Deserialize;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
/// This endpoint returns the user id
|
||||
pub struct GetStorageListRequest {
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct StorageItem {
|
||||
// NB: The path contains the user ID.
|
||||
pub bucket_pathname: Option<String>,
|
||||
|
||||
// NB: Other fields aren't necessary for us right now.
|
||||
}
|
||||
|
||||
pub async fn storage_list(req: GetStorageListRequest) -> Result<Vec<StorageItem>, MidjourneyError> {
|
||||
|
||||
let referer = format!("https://{}", req.hostname.as_str());
|
||||
|
||||
let url = format!("https://{}/api/storage-list", req.hostname.as_str());
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox139)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let cookie_header = req.cookie_header.trim();
|
||||
|
||||
if cookie_header.len() < 20 {
|
||||
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
|
||||
return Err(MidjourneyClientError::CookieTooShort.into());
|
||||
}
|
||||
|
||||
// NB: missing headers that were in the browser request:
|
||||
// -H 'sec-ch-ua-platform: "macOS"' \
|
||||
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
|
||||
|
||||
let mut http_request = client.get(url)
|
||||
.header("cookie", cookie_header)
|
||||
.header("Referer", &referer)
|
||||
.header("Referrer-Policy", "origin-when-cross-origin")
|
||||
.header("accept", "*/*")
|
||||
.header("accept-language", "en-US,en;q=0.8")
|
||||
.header("content-type", "application/json")
|
||||
.header("priority", "u=1, i")
|
||||
.header("sec-ch-ua-mobile", "?0")
|
||||
.header("sec-fetch-dest", "empty")
|
||||
.header("sec-fetch-mode", "cors")
|
||||
.header("sec-fetch-site", "same-origin")
|
||||
.header("x-csrf-protection", "1");
|
||||
|
||||
let http_request = http_request
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let response = client.execute(http_request)
|
||||
.await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
// if status == 301 {
|
||||
// for (name, value) in response.headers().iter() {
|
||||
// println!("- {}: {}", name.as_str(), value.to_str().unwrap());
|
||||
// }
|
||||
// }
|
||||
|
||||
let response_body = &response.text().await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
|
||||
return Err(MidjourneyApiError::CloudflareError(err).into());
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
[
|
||||
{
|
||||
"bucketPathname": "UUID/OTHER_HEX_ENTROPY.png",
|
||||
"shortUrl": null,
|
||||
"type": "image-prompt",
|
||||
"origin": "web",
|
||||
"timeCreated": TIMESTAMP,
|
||||
"hidden": false,
|
||||
"contentType": "image/png",
|
||||
"state": "loaded"
|
||||
}
|
||||
]
|
||||
*/
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawStorageItem {
|
||||
bucketPathname: Option<String>,
|
||||
}
|
||||
|
||||
let response : Vec<RawStorageItem> = serde_json::from_str(response_body)
|
||||
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
|
||||
|
||||
let items = response
|
||||
.into_iter()
|
||||
.map(|r| StorageItem {
|
||||
bucket_pathname: r.bucketPathname,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
use crate::endpoints::storage_list::{storage_list, GetStorageListRequest};
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
|
||||
let result = storage_list(GetStorageListRequest {
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}\n\n", result);
|
||||
|
||||
for provider in result {
|
||||
println!("Provider: {:?}", provider);
|
||||
}
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use log::{error, warn};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
pub struct SubmitJobRequest<'a> {
|
||||
pub prompt: &'a str,
|
||||
pub channel_id: &'a str,
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubmitJobResponse {
|
||||
/// On success, the job ID is returned.
|
||||
pub maybe_job_id: Option<String>,
|
||||
|
||||
/// On error, we have a list of error messages.
|
||||
pub maybe_errors: Option<Vec<SubmitJobError>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubmitJobError {
|
||||
pub error_type: Option<String>,
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse, MidjourneyError> {
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox139)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let referer = format!("https://{}", req.hostname.as_str());
|
||||
|
||||
// NB: Other recent clients use /api/app/submit-jobs, but this appears invalid.
|
||||
let url = format!("https://{}/api/submit-jobs", req.hostname.as_str());
|
||||
|
||||
let cookie_header = req.cookie_header.trim();
|
||||
|
||||
if cookie_header.len() < 20 {
|
||||
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
|
||||
return Err(MidjourneyClientError::CookieTooShort.into());
|
||||
}
|
||||
|
||||
let mut http_request = client.post(url)
|
||||
.header("cookie", cookie_header)
|
||||
.header("Referer", &referer)
|
||||
.header("Referrer-Policy", "origin-when-cross-origin")
|
||||
.header("accept", "*/*")
|
||||
.header("accept-language", "en-US,en;q=0.8")
|
||||
.header("content-type", "application/json")
|
||||
.header("priority", "u=1, i")
|
||||
.header("sec-ch-ua-mobile", "?0")
|
||||
.header("sec-fetch-dest", "empty")
|
||||
.header("sec-fetch-mode", "cors")
|
||||
.header("sec-fetch-site", "same-origin")
|
||||
.header("x-csrf-protection", "1");
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct F {
|
||||
mode: String,
|
||||
private: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[allow(non_snake_case)]
|
||||
struct Metadata {
|
||||
imagePrompts: u8,
|
||||
imageReferences: u8,
|
||||
characterReferences: u8,
|
||||
depthReferences: u8,
|
||||
lightboxOpen: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawRequest {
|
||||
f: F,
|
||||
channelId: String,
|
||||
metadata: Metadata,
|
||||
t: String,
|
||||
prompt: String,
|
||||
}
|
||||
|
||||
let body = RawRequest {
|
||||
f: F {
|
||||
mode: "fast".to_string(),
|
||||
private: false,
|
||||
},
|
||||
channelId: req.channel_id.to_string(),
|
||||
metadata: Metadata {
|
||||
imagePrompts: 0,
|
||||
imageReferences: 0,
|
||||
characterReferences: 0,
|
||||
depthReferences: 0,
|
||||
lightboxOpen: "".to_string(),
|
||||
},
|
||||
t: "imagine".to_string(),
|
||||
prompt: req.prompt.to_string(),
|
||||
};
|
||||
|
||||
let http_request = http_request
|
||||
.json(&body)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let response = client.execute(http_request)
|
||||
.await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
// if status == 301 {
|
||||
// for (name, value) in response.headers().iter() {
|
||||
// println!("- {}: {}", name.as_str(), value.to_str().unwrap());
|
||||
// }
|
||||
// }
|
||||
|
||||
let response_body = &response.text().await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
|
||||
return Err(MidjourneyApiError::CloudflareError(err).into());
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
{"success":[
|
||||
{
|
||||
"job_id":"UUID-like-string",
|
||||
"prompt":"prompt",
|
||||
"is_queued":false,
|
||||
"event_type":"diffusion",
|
||||
"flags":{
|
||||
"mode":"fast","visibility":"public"
|
||||
},
|
||||
"meta":{
|
||||
"height":1024,"width":1024,"batch_size":4,"parent_id":null,"parent_grid":null
|
||||
},
|
||||
"optimisticJobIndex":0,
|
||||
"personalization_codes":null
|
||||
}],
|
||||
"failure":[]}
|
||||
*/
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SuccessPayload {
|
||||
job_id: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct FailurePayload {
|
||||
r#type: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawResponse {
|
||||
success: Vec<SuccessPayload>,
|
||||
failure: Option<Vec<FailurePayload>>,
|
||||
}
|
||||
|
||||
let response = serde_json::from_str::<RawResponse>(response_body)
|
||||
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
|
||||
|
||||
let maybe_job_id = response.success
|
||||
.get(0)
|
||||
.map(|s| s.job_id.clone());
|
||||
|
||||
if maybe_job_id.is_none() {
|
||||
warn!("No job id found in body: {:?}", response_body);
|
||||
}
|
||||
|
||||
let mut maybe_errors = None;
|
||||
|
||||
if let Some(failures) = response.failure.as_ref() {
|
||||
if !failures.is_empty() {
|
||||
maybe_errors = Some(failures.iter()
|
||||
.map(|f| SubmitJobError {
|
||||
error_type: Some(f.r#type.clone()),
|
||||
message: Some(f.message.clone()),
|
||||
}).collect());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(SubmitJobResponse {
|
||||
maybe_job_id,
|
||||
maybe_errors,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::endpoints::submit_job::{submit_job, SubmitJobRequest};
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
|
||||
// Get channel id via:
|
||||
// https://identitytoolkit.googleapis.com/v1/accounts:lookup?key=[TOKEN]
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
let channel_id = read_to_trimmed_string("/Users/bt/secrets/midjourney/channel_id.txt")?;
|
||||
|
||||
let result = submit_job(SubmitJobRequest {
|
||||
prompt: "a modern n64 console",
|
||||
channel_id: &channel_id,
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}", result);
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
|
||||
use log::error;
|
||||
use serde::Deserialize;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
const USER_CONNECTIONS_URL :&str = "https://www.midjourney.com/api/user-connections";
|
||||
|
||||
pub struct GetUserConnectionsRequest {
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct UserData {
|
||||
pub provider_id: String,
|
||||
pub display_name: Option<String>,
|
||||
pub email: Option<String>,
|
||||
pub user_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Provider {
|
||||
pub provider_id: String,
|
||||
pub user_data: Option<UserData>,
|
||||
pub is_linked: bool,
|
||||
}
|
||||
|
||||
pub async fn user_connections(req: GetUserConnectionsRequest) -> Result<Vec<Provider>, MidjourneyError> {
|
||||
|
||||
let referer = format!("https://{}", req.hostname.as_str());
|
||||
|
||||
let url = format!("https://{}/api/user-connections", req.hostname.as_str());
|
||||
|
||||
let client = Client::builder()
|
||||
.emulation(Emulation::Firefox139)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let cookie_header = req.cookie_header.trim();
|
||||
|
||||
if cookie_header.len() < 20 {
|
||||
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
|
||||
return Err(MidjourneyClientError::CookieTooShort.into());
|
||||
}
|
||||
|
||||
// NB: missing headers that were in the browser request:
|
||||
// -H 'sec-ch-ua-platform: "macOS"' \
|
||||
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
|
||||
|
||||
let mut http_request = client.get(url)
|
||||
.header("cookie", cookie_header)
|
||||
.header("Referer", &referer)
|
||||
.header("Referrer-Policy", "origin-when-cross-origin")
|
||||
.header("accept", "*/*")
|
||||
.header("accept-language", "en-US,en;q=0.8")
|
||||
.header("content-type", "application/json")
|
||||
.header("priority", "u=1, i")
|
||||
.header("sec-ch-ua-mobile", "?0")
|
||||
.header("sec-fetch-dest", "empty")
|
||||
.header("sec-fetch-mode", "cors")
|
||||
.header("sec-fetch-site", "same-origin")
|
||||
.header("x-csrf-protection", "1");
|
||||
|
||||
let http_request = http_request
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let response = client.execute(http_request)
|
||||
.await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
// if status == 301 {
|
||||
// for (name, value) in response.headers().iter() {
|
||||
// println!("- {}: {}", name.as_str(), value.to_str().unwrap());
|
||||
// }
|
||||
// }
|
||||
|
||||
let response_body = &response.text().await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
|
||||
return Err(MidjourneyApiError::CloudflareError(err).into());
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
[
|
||||
{
|
||||
"providerId": "google.com",
|
||||
"userData": {
|
||||
"providerId": "google.com",
|
||||
"displayName": "John Doe",
|
||||
"email": "john.doe@gmail.com",
|
||||
"userId": "123456789000000001",
|
||||
},
|
||||
"isLinked": true
|
||||
},
|
||||
{
|
||||
"providerId": "discord.com",
|
||||
"isLinked": false
|
||||
}
|
||||
]
|
||||
*/
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawProviderUserData {
|
||||
providerId: String,
|
||||
displayName: Option<String>,
|
||||
email: Option<String>,
|
||||
userId: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawProvider {
|
||||
providerId: String,
|
||||
userData: Option<RawProviderUserData>,
|
||||
isLinked: bool,
|
||||
}
|
||||
|
||||
let response : Vec<RawProvider> = serde_json::from_str(response_body)
|
||||
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
|
||||
|
||||
let providers = response
|
||||
.into_iter()
|
||||
.map(|r| Provider {
|
||||
provider_id: r.providerId,
|
||||
user_data: r.userData.map(|u| UserData {
|
||||
provider_id: u.providerId,
|
||||
display_name: u.displayName,
|
||||
email: u.email,
|
||||
user_id: u.userId,
|
||||
}),
|
||||
is_linked: r.isLinked,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(providers)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::endpoints::user_connections::{user_connections, GetUserConnectionsRequest};
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
|
||||
let result = user_connections(GetUserConnectionsRequest {
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}\n\n", result);
|
||||
|
||||
for provider in result {
|
||||
println!("Provider: {:?}", provider);
|
||||
}
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
use cloudflare_errors::cloudflare_error::CloudflareError;
|
||||
use errors::AnyhowError;
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::io;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum MidjourneyApiError {
|
||||
/// There was no user ID in an otherwise valid-looking response.
|
||||
NoUserId,
|
||||
|
||||
/// No user props in index HTML payload.
|
||||
/// We need this for the user ID and the websocket token.
|
||||
NoUserProps,
|
||||
|
||||
/// No initialAuthUser in the index HTML payload.
|
||||
/// We need this for the user ID and the websocket token.
|
||||
NoInitialAuthUser,
|
||||
|
||||
/// 400. The request was invalid.
|
||||
InvalidRequest(String),
|
||||
|
||||
/// 401. The request was not authorized.
|
||||
Unauthorized(String),
|
||||
|
||||
/// 403. The request was forbidden.
|
||||
Forbidden(String),
|
||||
|
||||
/// 404. The requested resource was not found.
|
||||
NotFound(String),
|
||||
|
||||
/// 429. Too many requests.
|
||||
TooManyRequests(String),
|
||||
|
||||
/// 500. An internal server error occurred.
|
||||
InternalServerError {
|
||||
body: String,
|
||||
backend_hostname: Option<String>,
|
||||
},
|
||||
|
||||
/// Eg. when downloading images
|
||||
UnknownHttpFailure {
|
||||
status_code: u16,
|
||||
body: String,
|
||||
},
|
||||
|
||||
/// Cloudflare errors.
|
||||
CloudflareError(CloudflareError),
|
||||
|
||||
/// A deserialization error with the response.
|
||||
DeserializationError(serde_json::Error),
|
||||
|
||||
/// The request timed out.
|
||||
Timeout(String),
|
||||
|
||||
/// A network error occurred.
|
||||
NetworkError(String),
|
||||
|
||||
/// An error doing file I/O (on our side)
|
||||
IoError(io::Error),
|
||||
|
||||
/// Another type of error.
|
||||
Other(AnyhowError),
|
||||
}
|
||||
|
||||
impl Error for MidjourneyApiError {}
|
||||
|
||||
impl Display for MidjourneyApiError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
// Response body errors
|
||||
Self::NoUserId => write!(f, "No user ID found in the response body."),
|
||||
Self::NoUserProps => write!(f, "No user properties found in the index HTML payload."),
|
||||
Self::NoInitialAuthUser => write!(f, "No initialAuthUser in the index HTML payload."),
|
||||
// Server response code errors
|
||||
Self::InvalidRequest(msg) => write!(f, "Invalid request: {}", msg),
|
||||
Self::Unauthorized(msg) => write!(f, "Unauthorized: {}", msg),
|
||||
Self::Forbidden(msg) => write!(f, "Forbidden: {}", msg),
|
||||
Self::NotFound(msg) => write!(f, "Not found: {}", msg),
|
||||
Self::TooManyRequests(msg) => write!(f, "Too many requests: {}", msg),
|
||||
Self::InternalServerError {body, backend_hostname} =>
|
||||
write!(f, "Internal Server Error; backend hostname: {:?} ; body: {}; ", backend_hostname, body),
|
||||
Self::UnknownHttpFailure {status_code, body} =>
|
||||
write!(f, "Unknown HTTP failure; status code: {}; body: {}", status_code, body),
|
||||
// Deserialization errors
|
||||
// Server response handling errors
|
||||
Self::DeserializationError(error) => write!(f, "Deserialization error: {}", error),
|
||||
// Network errors
|
||||
Self::Timeout(msg) => write!(f, "Timeout: {}", msg),
|
||||
Self::NetworkError(msg) => write!(f, "Network error: {}", msg),
|
||||
// Cloudflare errors
|
||||
Self::CloudflareError(error) => write!(f, "Cloudflare Error: {}", error),
|
||||
// I/O errors
|
||||
Self::IoError(error) => write!(f, "IO error: {}", error),
|
||||
// Other
|
||||
Self::Other(error) => write!(f, "Other error: {}", error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for MidjourneyApiError {
|
||||
fn from(error: serde_json::Error) -> Self {
|
||||
Self::DeserializationError(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<io::Error> for MidjourneyApiError {
|
||||
fn from(error: io::Error) -> Self {
|
||||
Self::IoError(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<CloudflareError> for MidjourneyApiError {
|
||||
fn from(error: CloudflareError) -> Self {
|
||||
Self::CloudflareError(error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
use std::error::Error;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// An error arising from the client, either before or after the request is sent.
|
||||
#[derive(Debug)]
|
||||
pub enum MidjourneyClientError {
|
||||
/// The cookie was too short to be valid.
|
||||
CookieTooShort,
|
||||
|
||||
/// Tried to look up an image out of range [0,4)
|
||||
InvalidImageIndex,
|
||||
|
||||
/// Could not determine the filetype for the file.
|
||||
FileTypeNotKnown(PathBuf),
|
||||
|
||||
/// Could not handle the file type.
|
||||
FileTypeNotHandled(PathBuf),
|
||||
|
||||
/// An error doing file I/O (on our side)
|
||||
IoError(io::Error),
|
||||
|
||||
// /// An error setting up the client
|
||||
// ReqwestError(reqwest::Error),
|
||||
WreqError(wreq::Error),
|
||||
|
||||
/// A serialization error with the request.
|
||||
SerializationError(serde_json::Error),
|
||||
}
|
||||
|
||||
impl Error for MidjourneyClientError {}
|
||||
|
||||
impl Display for MidjourneyClientError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::CookieTooShort => write!(f, "The cookie header is too short to be valid."),
|
||||
Self::InvalidImageIndex => write!(f, "The image index is out of range. Valid range is [0, 4)."),
|
||||
Self::FileTypeNotKnown(path) => write!(f, "Could not determine the filetype for the file: {:?}", path),
|
||||
Self::FileTypeNotHandled(path) => write!(f, "Could not handle the file type for the file: {:?}", path),
|
||||
Self::IoError(err) => write!(f, "IO error: {}", err),
|
||||
//Self::ReqwestError(err) => write!(f, "Reqwest client error: {}", err),
|
||||
Self::WreqError(err) => write!(f, "Wreq client error: {}", err),
|
||||
Self::SerializationError(err) => write!(f, "Request serialization error: {}", err),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<io::Error> for MidjourneyClientError {
|
||||
fn from(error: io::Error) -> Self {
|
||||
Self::IoError(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for MidjourneyClientError {
|
||||
fn from(error: serde_json::Error) -> Self {
|
||||
Self::SerializationError(error)
|
||||
}
|
||||
}
|
||||
|
||||
/*impl From<reqwest::Error> for MidjourneyClientError {
|
||||
fn from(error: reqwest::Error) -> Self {
|
||||
Self::ReqwestError(error)
|
||||
}
|
||||
}*/
|
||||
|
||||
impl From<wreq::Error> for MidjourneyClientError {
|
||||
fn from(error: wreq::Error) -> Self {
|
||||
Self::WreqError(error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use std::error::Error;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum MidjourneyError {
|
||||
Api(MidjourneyApiError),
|
||||
Client(MidjourneyClientError),
|
||||
}
|
||||
|
||||
impl Error for MidjourneyError {}
|
||||
|
||||
impl std::fmt::Display for MidjourneyError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Api(e) => write!(f, "Midjourney API error: {:?}", e),
|
||||
Self::Client(e) => write!(f, "Midjourney Client error: {:?}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<MidjourneyApiError> for MidjourneyError {
|
||||
fn from(error: MidjourneyApiError) -> Self {
|
||||
Self::Api(error)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<MidjourneyClientError> for MidjourneyError {
|
||||
fn from(error: MidjourneyClientError) -> Self {
|
||||
Self::Client(error)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod midjourney_error;
|
||||
pub mod midjourney_api_error;
|
||||
pub mod midjourney_client_error;
|
||||
@@ -0,0 +1,8 @@
|
||||
//! A midjourney API client.
|
||||
|
||||
pub mod client;
|
||||
pub mod credentials;
|
||||
pub mod endpoints;
|
||||
pub mod error;
|
||||
pub mod recipes;
|
||||
pub mod utils;
|
||||
@@ -0,0 +1,20 @@
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ChannelId {
|
||||
/// Channel based on the user ID.
|
||||
/// This is what Midjourney uses by default.
|
||||
UserId(MidjourneyUserId),
|
||||
|
||||
// Bare string ID
|
||||
Raw(String),
|
||||
}
|
||||
|
||||
impl ChannelId {
|
||||
pub fn to_string(&self) -> String {
|
||||
match self {
|
||||
Self::UserId(user_id) => format!("singleplayer_{}", user_id.0),
|
||||
Self::Raw(id) => id.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use crate::endpoints::storage_list::{storage_list, GetStorageListRequest};
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
|
||||
pub struct GetMidjourneyUserIdRequest {
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
pub async fn get_midjourney_user_id_from_storage(req: GetMidjourneyUserIdRequest) -> Result<MidjourneyUserId, MidjourneyError> {
|
||||
let items = storage_list(GetStorageListRequest {
|
||||
hostname: req.hostname,
|
||||
cookie_header: req.cookie_header,
|
||||
}).await?;
|
||||
|
||||
let user_id = items
|
||||
.into_iter()
|
||||
.find_map(|item| item.bucket_pathname)
|
||||
.and_then(|path| path.split('/').next().map(|id| id.to_owned()))
|
||||
.map(|id| MidjourneyUserId(id.to_string()))
|
||||
.ok_or(MidjourneyApiError::NoUserId)?;
|
||||
|
||||
Ok(user_id)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::recipes::get_midjourney_user_id_from_storage::{get_midjourney_user_id_from_storage, GetMidjourneyUserIdRequest};
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
|
||||
let result = get_midjourney_user_id_from_storage(GetMidjourneyUserIdRequest {
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}\n\n", result);
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageRequest};
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use dom_query::Document;
|
||||
use serde::Deserialize;
|
||||
|
||||
pub struct GetUserInfoRequest {
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GetUserInfoResponse {
|
||||
pub user_id: Option<MidjourneyUserId>,
|
||||
pub email: Option<String>,
|
||||
pub websocket_token: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn get_user_info(req: GetUserInfoRequest) -> Result<GetUserInfoResponse, MidjourneyError> {
|
||||
let index_html = get_index_page_html(GetIndexPageRequest {
|
||||
hostname: req.hostname,
|
||||
cookie_header: req.cookie_header,
|
||||
}).await?;
|
||||
|
||||
/*
|
||||
<script id="initialProps" type="application/json">
|
||||
{
|
||||
"initialAuthUser": {
|
||||
"idpID": "ENTROPY",
|
||||
"id": "UUID",
|
||||
"displayName": "uNUMBER",
|
||||
"email": "USER@gmail.com",
|
||||
"emailVerified": true,
|
||||
"photoURL": "https://lh3.googleusercontent.com/a/ENTROPY",
|
||||
"abilities": {
|
||||
"admin": false,
|
||||
"developer": false,
|
||||
"accepted_tos": true,
|
||||
"moderator": false,
|
||||
"guide": false,
|
||||
"vip": false,
|
||||
"employee": false,
|
||||
"content_moderator": false,
|
||||
"gdpr_preferences_set": false,
|
||||
"gdpr_allow_advertising": false,
|
||||
"gdpr_allow_analytics": false,
|
||||
"editor_whitelist": false,
|
||||
"allow_nsfw": false,
|
||||
"cooldowns_removed": false,
|
||||
"blocked": false,
|
||||
"community": false,
|
||||
"billing": false,
|
||||
"system_verified": false,
|
||||
"simplified_flow": false,
|
||||
"total_jobs": 57,
|
||||
"is_sub_at_least_one_year_old": false,
|
||||
"can_private": false,
|
||||
"can_relax": false,
|
||||
"has_discord": false,
|
||||
"delete_at": null,
|
||||
"subscription": {
|
||||
"type": "subscription",
|
||||
"plan_type": "basic",
|
||||
"recurring": "month"
|
||||
},
|
||||
"relax_video_job_concurrency": 0
|
||||
},
|
||||
"websocketToken": "WEBSOCKET_TOKEN_ENTROPY",
|
||||
"providers": [
|
||||
"google.com"
|
||||
],
|
||||
"system": "firebase",
|
||||
"streamToken": "STREAM_TOKEN_ENTROPY"
|
||||
},
|
||||
"urlFromInitialServerRequest": "/"
|
||||
}
|
||||
</script>
|
||||
*/
|
||||
|
||||
let document = Document::from(index_html);
|
||||
let maybe_nodes = document.select("script[id=initialProps]");
|
||||
|
||||
let node = match maybe_nodes.get(0) {
|
||||
Some(node) => node,
|
||||
None => {
|
||||
return Err(MidjourneyApiError::NoUserProps.into());
|
||||
}
|
||||
};
|
||||
|
||||
let inner_json = node.inner_html().to_string();
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawAuthUser {
|
||||
// NB: The Midjourney user id.
|
||||
id: Option<String>,
|
||||
email: Option<String>,
|
||||
websocketToken: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct RawBody {
|
||||
initialAuthUser: Option<RawAuthUser>,
|
||||
}
|
||||
|
||||
let response : RawBody = serde_json::from_str(&inner_json)
|
||||
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
|
||||
|
||||
let props = match response.initialAuthUser {
|
||||
Some(props) => props,
|
||||
None => {
|
||||
return Err(MidjourneyApiError::NoInitialAuthUser.into());
|
||||
}
|
||||
};
|
||||
|
||||
Ok(GetUserInfoResponse {
|
||||
user_id: props.id
|
||||
.map(|id| MidjourneyUserId::from_string(id)),
|
||||
email: props.email
|
||||
.map(|email| email.to_string()),
|
||||
websocket_token: props.websocketToken
|
||||
.map(|email| email.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::recipes::get_user_info::{get_user_info, GetUserInfoRequest};
|
||||
use errors::AnyhowResult;
|
||||
use filesys::read_to_trimmed_string::read_to_trimmed_string;
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn test() -> AnyhowResult<()> {
|
||||
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
|
||||
|
||||
let result = get_user_info(GetUserInfoRequest {
|
||||
cookie_header,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
}).await?;
|
||||
|
||||
println!("Response: {:?}\n\n", result);
|
||||
|
||||
assert_eq!(1, 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
mod get_midjourney_user_id_from_storage;
|
||||
pub mod channel_id;
|
||||
pub mod get_user_info;
|
||||
pub mod text_to_image;
|
||||
@@ -0,0 +1,49 @@
|
||||
use crate::client::midjourney_hostname::MidjourneyHostname;
|
||||
use crate::endpoints::submit_job::{submit_job, SubmitJobRequest};
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use crate::recipes::channel_id::ChannelId;
|
||||
|
||||
pub struct TextToImageRequest<'a> {
|
||||
pub prompt: &'a str,
|
||||
pub channel_id: &'a ChannelId,
|
||||
pub hostname: MidjourneyHostname,
|
||||
pub cookie_header: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TextToImageResponse {
|
||||
/// On success, the job ID is returned.
|
||||
pub maybe_job_id: Option<String>,
|
||||
|
||||
/// On error, we have a list of error messages.
|
||||
pub maybe_errors: Option<Vec<TextToImageError>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TextToImageError {
|
||||
pub error_type: Option<String>,
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
/// Slightly more ergonomic text-to-image API.
|
||||
/// As we add more `submit_job()` cases, we'll keep this simple.
|
||||
pub async fn text_to_image(req: TextToImageRequest<'_>) -> Result<TextToImageResponse, MidjourneyError> {
|
||||
let channel_id = req.channel_id.to_string();
|
||||
|
||||
let response = submit_job(SubmitJobRequest {
|
||||
prompt: req.prompt,
|
||||
channel_id: &channel_id,
|
||||
hostname: req.hostname,
|
||||
cookie_header: req.cookie_header,
|
||||
}).await?;
|
||||
|
||||
Ok(TextToImageResponse {
|
||||
maybe_job_id: response.maybe_job_id,
|
||||
maybe_errors: response.maybe_errors.map(|errs| {
|
||||
errs.into_iter().map(|e| TextToImageError {
|
||||
error_type: e.error_type,
|
||||
message: e.message,
|
||||
}).collect()
|
||||
}),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
|
||||
pub fn get_image_url(job_id: &str, image_index: u8) -> Result<String, MidjourneyError> {
|
||||
if image_index > 3 {
|
||||
return Err(MidjourneyClientError::InvalidImageIndex.into())
|
||||
}
|
||||
Ok(format!("https://cdn.midjourney.com/{}/0_{}.png", job_id, image_index))
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
use crate::error::midjourney_api_error::MidjourneyApiError;
|
||||
use crate::error::midjourney_client_error::MidjourneyClientError;
|
||||
use crate::error::midjourney_error::MidjourneyError;
|
||||
use crate::utils::get_image_url::get_image_url;
|
||||
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
|
||||
use wreq::Client;
|
||||
use wreq_util::Emulation;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ImageDownloaderClient {
|
||||
client: Client,
|
||||
}
|
||||
|
||||
impl ImageDownloaderClient {
|
||||
pub fn create() -> Result<Self, MidjourneyClientError> {
|
||||
Ok(Self {
|
||||
client: Client::builder()
|
||||
.emulation(Emulation::Firefox139)
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn download_image(&self, job_id: &str, image_index: u8) -> anyhow::Result<Vec<u8>, MidjourneyError> {
|
||||
let url = get_image_url(job_id, image_index)?;
|
||||
|
||||
// TODO: Cookies
|
||||
// TODO: Cache control headers?
|
||||
let mut http_request = self.client.get(url)
|
||||
//.header("Referrer-Policy", "origin-when-cross-origin")
|
||||
//.header("content-type", "application/json")
|
||||
//.header("sec-fetch-user", "?1")
|
||||
//.header("upgrade-insecure-requests", "1");
|
||||
.header("Referrer", "https://www.midjourney.com/")
|
||||
.header("accept", "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8")
|
||||
.header("accept-language", "en-US,en;q=0.8")
|
||||
.header("priority", "i")
|
||||
.header("sec-ch-ua-mobile", "?0")
|
||||
.header("sec-fetch-dest", "image")
|
||||
.header("sec-fetch-mode", "no-cors")
|
||||
.header("sec-fetch-site", "same-site");
|
||||
|
||||
let http_request = http_request
|
||||
.build()
|
||||
.map_err(|err| MidjourneyClientError::WreqError(err))?;
|
||||
|
||||
let response = self.client.execute(http_request)
|
||||
.await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
let response_bytes = response.bytes().await
|
||||
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
let response_body = String::from_utf8_lossy(&response_bytes).to_string();
|
||||
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
|
||||
return Err(MidjourneyApiError::CloudflareError(err).into());
|
||||
}
|
||||
|
||||
return Err(MidjourneyApiError::UnknownHttpFailure {
|
||||
status_code: status.as_u16(),
|
||||
body: response_body,
|
||||
}.into());
|
||||
}
|
||||
|
||||
Ok(response_bytes.to_vec())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod get_image_url;
|
||||
pub mod image_downloader_client;
|
||||
@@ -59,7 +59,7 @@ impl Display for ApiError {
|
||||
ApiError::Unauthorized(msg) => write!(f, "Unauthorized: {}", msg),
|
||||
ApiError::Forbidden(msg) => write!(f, "Forbidden: {}", msg),
|
||||
ApiError::NotFound(msg) => write!(f, "Not found: {}", msg),
|
||||
ApiError::TooManyRequests(msg) => write!(f, "Too many requests: {}", msg),
|
||||
ApiError::TooManyRequests(msg) => write!(f, "Too many requests to Storyteller backend: {}", msg),
|
||||
ApiError::InternalServerError {body, backend_hostname} =>
|
||||
write!(f, "Internal Server Error; backend hostname: {:?} ; body: {}; ", backend_hostname, body),
|
||||
// Server response handling errors
|
||||
|
||||
@@ -7,6 +7,9 @@ use crate::error::api_error::ApiError;
|
||||
/// An error arising from the client, either before or after the request is sent.
|
||||
#[derive(Debug)]
|
||||
pub enum ClientError {
|
||||
/// An invalid request was detected before it was sent.
|
||||
InvalidPreflightRequest(String),
|
||||
|
||||
/// Could not determine the filetype for the file.
|
||||
FileTypeNotKnown(PathBuf),
|
||||
|
||||
@@ -28,6 +31,7 @@ impl Error for ClientError {}
|
||||
impl Display for ClientError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::InvalidPreflightRequest(msg) => write!(f, "Invalid preflight request: {}", msg),
|
||||
Self::FileTypeNotKnown(path) => write!(f, "Could not determine the filetype for the file: {:?}", path),
|
||||
Self::FileTypeNotHandled(path) => write!(f, "Could not handle the file type for the file: {:?}", path),
|
||||
Self::IoError(err) => write!(f, "IO error: {}", err),
|
||||
|
||||
@@ -13,8 +13,8 @@ impl Error for StorytellerError {}
|
||||
impl std::fmt::Display for StorytellerError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
StorytellerError::Api(e) => write!(f, "API error: {:?}", e),
|
||||
StorytellerError::Client(e) => write!(f, "Client error: {:?}", e),
|
||||
StorytellerError::Api(e) => write!(f, "Storyteller API error: {:?}", e),
|
||||
StorytellerError::Client(e) => write!(f, "Storyteller Client error: {:?}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ pub mod error;
|
||||
pub mod generate;
|
||||
pub mod jobs;
|
||||
pub mod media_files;
|
||||
pub mod prompts;
|
||||
pub mod recipes;
|
||||
pub mod shared_response_types;
|
||||
pub mod utils;
|
||||
|
||||
+27
@@ -16,6 +16,7 @@ use enums::by_table::media_files::media_file_class::MediaFileClass;
|
||||
use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory;
|
||||
use enums::by_table::media_files::media_file_type::MediaFileType;
|
||||
use enums::by_table::model_weights::weights_category::WeightsCategory;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::visibility::Visibility;
|
||||
use enums::no_table::style_transfer::style_transfer_name::StyleTransferName;
|
||||
use enums_public::by_table::model_weights::public_weights_types::PublicWeightsType;
|
||||
@@ -42,6 +43,14 @@ pub struct UploadImageFromFileArgs<'a, P: AsRef<Path>> {
|
||||
|
||||
/// If true, we should hide the image from the user's gallery.
|
||||
pub is_intermediate_system_file: bool,
|
||||
|
||||
/// If provided, this is the prompt that this image is associated with.
|
||||
/// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
|
||||
pub maybe_prompt_token: Option<&'a PromptToken>,
|
||||
|
||||
// /// If provided, this is the service provider that created the image.
|
||||
// /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set.
|
||||
// pub maybe_generation_provider: Option<GenerationProvider>,
|
||||
}
|
||||
|
||||
|
||||
@@ -49,6 +58,8 @@ pub struct UploadImageFromFileArgs<'a, P: AsRef<Path>> {
|
||||
pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
|
||||
args: UploadImageFromFileArgs<'_, P>,
|
||||
) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
|
||||
|
||||
validate_args(&args)?;
|
||||
|
||||
let url = get_route(args.api_host);
|
||||
|
||||
@@ -70,6 +81,10 @@ pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
|
||||
if args.is_intermediate_system_file {
|
||||
form = form.text("is_intermediate_system_file", "true");
|
||||
}
|
||||
|
||||
if let Some(prompt_token) = &args.maybe_prompt_token {
|
||||
form = form.text("maybe_prompt_token", prompt_token.to_string());
|
||||
}
|
||||
|
||||
let mut request_builder = client.post(url)
|
||||
.header("User-Agent", USER_AGENT)
|
||||
@@ -109,3 +124,15 @@ pub struct UploadImageMediaFileSuccessResponse {
|
||||
pub success: bool,
|
||||
pub media_file_token: MediaFileToken,
|
||||
}
|
||||
|
||||
fn validate_args<P: AsRef<Path>>(
|
||||
args: &UploadImageFromFileArgs<'_, P>,
|
||||
) -> Result<(), StorytellerError> {
|
||||
if args.is_intermediate_system_file && args.maybe_prompt_token.is_some() {
|
||||
return Err(StorytellerError::Client(ClientError::InvalidPreflightRequest(
|
||||
"Cannot set `is_intermediate_system_file` to true if `maybe_prompt_token` is provided."
|
||||
.to_string())));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use crate::error::storyteller_error::StorytellerError;
|
||||
use crate::utils::api_host::ApiHost;
|
||||
use crate::utils::basic_json_post_request::basic_json_post_request;
|
||||
use artcraft_api_defs::prompts::create_prompt::{CreatePromptRequest, CreatePromptResponse, CREATE_PROMPT_PATH};
|
||||
|
||||
pub async fn create_prompt(
|
||||
api_host: &ApiHost,
|
||||
maybe_creds: Option<&StorytellerCredentialSet>,
|
||||
request: CreatePromptRequest,
|
||||
) -> Result<CreatePromptResponse, StorytellerError> {
|
||||
Ok(basic_json_post_request(
|
||||
api_host,
|
||||
CREATE_PROMPT_PATH,
|
||||
maybe_creds,
|
||||
request,
|
||||
).await?)
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod create_prompt;
|
||||
@@ -10,6 +10,7 @@ use mimetypes::mimetype_info::mimetype_info::MimetypeInfo;
|
||||
use serde_derive::Deserialize;
|
||||
use std::path::Path;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
use tokens::tokens::prompts::PromptToken;
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
pub struct UploadMediaFileSuccessResponse {
|
||||
@@ -22,6 +23,7 @@ pub async fn upload_media_file_from_file<P: AsRef<Path>>(
|
||||
api_host: &ApiHost,
|
||||
maybe_creds: Option<&StorytellerCredentialSet>,
|
||||
path: P,
|
||||
maybe_prompt_token: Option<&PromptToken>,
|
||||
) -> Result<UploadMediaFileSuccessResponse, StorytellerError> {
|
||||
|
||||
let maybe_type = MimetypeInfo::get_for_path(path.as_ref())
|
||||
@@ -54,6 +56,7 @@ pub async fn upload_media_file_from_file<P: AsRef<Path>>(
|
||||
maybe_creds,
|
||||
path,
|
||||
is_intermediate_system_file: false,
|
||||
maybe_prompt_token,
|
||||
}).await;
|
||||
match result {
|
||||
Ok(result) => Ok(UploadMediaFileSuccessResponse {
|
||||
|
||||
@@ -23,25 +23,27 @@ tauri-build = { version = "2.1.1", features = [] }
|
||||
|
||||
[dependencies]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration = { workspace = true }
|
||||
workspace-build-acceleration.workspace = true
|
||||
|
||||
# Internal
|
||||
artcraft_api_defs = { workspace = true }
|
||||
build_metadata = { workspace = true }
|
||||
enums = { workspace = true }
|
||||
errors = { workspace = true }
|
||||
fal_client = { workspace = true }
|
||||
filesys = { workspace = true }
|
||||
idempotency = { workspace = true }
|
||||
images = { workspace = true }
|
||||
memory_store = { workspace = true }
|
||||
mimetypes = { workspace = true }
|
||||
openai_sora_client = { workspace = true }
|
||||
primitives = { workspace = true }
|
||||
artcraft_api_defs.workspace = true
|
||||
build_metadata.workspace = true
|
||||
cookie_store.workspace = true
|
||||
enums.workspace = true
|
||||
errors.workspace = true
|
||||
fal_client.workspace = true
|
||||
filesys.workspace = true
|
||||
idempotency.workspace = true
|
||||
images.workspace = true
|
||||
memory_store.workspace = true
|
||||
midjourney_client.workspace = true
|
||||
mimetypes.workspace = true
|
||||
openai_sora_client.workspace = true
|
||||
primitives.workspace = true
|
||||
sqlite_tasks.workspace = true
|
||||
storyteller_client = { workspace = true }
|
||||
tokens = { workspace = true }
|
||||
web_base64 = { workspace = true }
|
||||
storyteller_client.workspace = true
|
||||
tokens.workspace = true
|
||||
web_base64.workspace = true
|
||||
|
||||
# External
|
||||
anyhow = { workspace = true }
|
||||
|
||||
+15
@@ -3,6 +3,7 @@ use crate::core::commands::enqueue::image::generic::handle_image_fal::handle_ima
|
||||
use crate::core::commands::enqueue::image::gpt_image_1::handle_gpt_image_1::handle_gpt_image_1;
|
||||
use crate::core::commands::enqueue::image::gpt_image_1::handle_gpt_image_1_sora::handle_gpt_image_1_sora;
|
||||
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
|
||||
use crate::core::commands::enqueue::image::midjourney::handle_midjourney::handle_midjourney;
|
||||
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
|
||||
use crate::core::commands::response::failure_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus};
|
||||
use crate::core::commands::response::shorthand::Response;
|
||||
@@ -18,6 +19,7 @@ use crate::core::state::provider_priority::{Provider, ProviderPriorityStore};
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
|
||||
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
@@ -93,6 +95,7 @@ pub async fn enqueue_text_to_image_command(
|
||||
provider_priority_store: State<'_, ProviderPriorityStore>,
|
||||
task_database: State<'_, TaskDatabase>,
|
||||
fal_creds_manager: State<'_, FalCredentialManager>,
|
||||
mj_creds_manager: State<'_, MidjourneyCredentialManager>,
|
||||
storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
|
||||
fal_task_queue: State<'_, FalTaskQueue>,
|
||||
sora_creds_manager: State<'_, SoraCredentialManager>,
|
||||
@@ -108,6 +111,7 @@ pub async fn enqueue_text_to_image_command(
|
||||
&provider_priority_store,
|
||||
&task_database,
|
||||
&fal_creds_manager,
|
||||
&mj_creds_manager,
|
||||
&storyteller_creds_manager,
|
||||
&app_env_configs,
|
||||
&fal_task_queue,
|
||||
@@ -173,6 +177,7 @@ pub async fn handle_request(
|
||||
provider_priority_store: &ProviderPriorityStore,
|
||||
task_database: &TaskDatabase,
|
||||
fal_creds_manager: &FalCredentialManager,
|
||||
mj_creds_manager: &MidjourneyCredentialManager,
|
||||
storyteller_creds_manager: &StorytellerCredentialManager,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
fal_task_queue: &FalTaskQueue,
|
||||
@@ -189,6 +194,7 @@ pub async fn handle_request(
|
||||
&storyteller_creds_manager,
|
||||
&app_env_configs,
|
||||
&fal_task_queue,
|
||||
&mj_creds_manager,
|
||||
&sora_creds_manager,
|
||||
&sora_task_queue,
|
||||
).await;
|
||||
@@ -219,6 +225,7 @@ pub async fn dispatch_request(
|
||||
storyteller_creds_manager: &StorytellerCredentialManager,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
fal_task_queue: &FalTaskQueue,
|
||||
mj_creds_manager: &MidjourneyCredentialManager,
|
||||
sora_creds_manager: &SoraCredentialManager,
|
||||
sora_task_queue: &SoraTaskQueue,
|
||||
) -> Result<TaskEnqueueSuccess, InternalImageError> {
|
||||
@@ -239,6 +246,14 @@ pub async fn dispatch_request(
|
||||
sora_task_queue,
|
||||
).await;
|
||||
}
|
||||
Some(ImageModel::Midjourney) => {
|
||||
return handle_midjourney(
|
||||
app,
|
||||
request,
|
||||
app_env_configs,
|
||||
mj_creds_manager,
|
||||
).await;
|
||||
}
|
||||
_ => {
|
||||
// Fall-through
|
||||
}
|
||||
|
||||
+5
-2
@@ -57,8 +57,11 @@ pub async fn handle_image_artcraft(
|
||||
None => {
|
||||
return Err(InternalImageError::NoModelSpecified);
|
||||
}
|
||||
Some(ImageModel::GptImage1) => {
|
||||
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: an earlier branch should have handled gpt-image-1 images for Artcraft")));
|
||||
Some(
|
||||
ImageModel::GptImage1 |
|
||||
ImageModel::Midjourney
|
||||
) => {
|
||||
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: another branch should handle this: {:?}", request.model)));
|
||||
}
|
||||
Some(ImageModel::Recraft3) => {
|
||||
return Err(InternalImageError::AnyhowError(anyhow!("not yet implemented in Artcraft")));
|
||||
|
||||
@@ -51,8 +51,11 @@ pub async fn handle_image_fal(
|
||||
None => {
|
||||
return Err(InternalImageError::NoModelSpecified);
|
||||
}
|
||||
Some(ImageModel::GptImage1) => {
|
||||
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: fal is handling sora images")));
|
||||
Some(
|
||||
ImageModel::GptImage1 |
|
||||
ImageModel::Midjourney
|
||||
) => {
|
||||
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: another branch should handle this: {:?}", request.model)));
|
||||
}
|
||||
Some(
|
||||
ImageModel::Flux1Dev |
|
||||
|
||||
@@ -2,6 +2,7 @@ use crate::core::artcraft_error::ArtcraftError;
|
||||
use base64::DecodeError;
|
||||
use errors::AnyhowError;
|
||||
use fal_client::error::fal_error_plus::FalErrorPlus;
|
||||
use midjourney_client::error::midjourney_error::MidjourneyError;
|
||||
use openai_sora_client::sora_error::SoraError;
|
||||
use storyteller_client::error::storyteller_error::StorytellerError;
|
||||
|
||||
@@ -10,9 +11,18 @@ pub enum InternalImageError {
|
||||
NoModelSpecified,
|
||||
NoProviderAvailable,
|
||||
NeedsFalApiKey,
|
||||
|
||||
NeedsMidjourneyCredentials,
|
||||
NeedsMidjourneyUserInfo,
|
||||
NeedsMidjourneyUserId,
|
||||
|
||||
/// Generic - possibly this came with an error message that we sent as an event.
|
||||
MidjourneyJobEnqueueFailed,
|
||||
|
||||
NeedsSoraCredentials,
|
||||
NeedsStorytellerCredentials,
|
||||
FalError(FalErrorPlus),
|
||||
MidjourneyError(MidjourneyError),
|
||||
SoraError(SoraError),
|
||||
AnyhowError(AnyhowError),
|
||||
StorytellerError(StorytellerError),
|
||||
|
||||
+142
@@ -0,0 +1,142 @@
|
||||
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize};
|
||||
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
|
||||
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize};
|
||||
use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError;
|
||||
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
|
||||
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
|
||||
use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent;
|
||||
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
|
||||
use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
|
||||
use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent;
|
||||
use crate::core::model::image_models::ImageModel;
|
||||
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::state::provider_priority::ProviderPriorityStore;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use artcraft_api_defs::generate::image::generate_gpt_image_1_text_to_image::{GenerateGptImage1TextToImageImageQuality, GenerateGptImage1TextToImageImageSize, GenerateGptImage1TextToImageNumImages, GenerateGptImage1TextToImageRequest};
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::tauri::tasks::task_type::TaskType;
|
||||
use fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs};
|
||||
use idempotency::uuid::generate_random_uuid;
|
||||
use log::{error, info};
|
||||
use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
|
||||
use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest};
|
||||
use midjourney_client::error::midjourney_api_error::MidjourneyApiError;
|
||||
use midjourney_client::recipes::channel_id::ChannelId;
|
||||
use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageError, TextToImageRequest};
|
||||
use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
|
||||
use openai_sora_client::recipes::simple_image_gen_with_session_auto_renew::{simple_image_gen_with_session_auto_renew, SimpleImageGenAutoRenewRequest};
|
||||
use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages};
|
||||
use std::time::Duration;
|
||||
use storyteller_client::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image;
|
||||
use storyteller_client::generate::image::generate_gpt_image_1_text_to_image::generate_gpt_image_1_text_to_image;
|
||||
use tauri::AppHandle;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
pub async fn handle_midjourney(
|
||||
app: &AppHandle,
|
||||
request: EnqueueTextToImageRequest,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
mj_creds_manager: &MidjourneyCredentialManager,
|
||||
) -> Result<TaskEnqueueSuccess, InternalImageError> {
|
||||
|
||||
let creds = match mj_creds_manager.maybe_copy_cookie_store() {
|
||||
Ok(Some(creds)) => creds,
|
||||
Ok(None) => {
|
||||
return Err(InternalImageError::NeedsMidjourneyCredentials);
|
||||
}
|
||||
Err(err) => {
|
||||
error!("Error reading Midjourney credentials: {:?}", err);
|
||||
return Err(InternalImageError::NeedsMidjourneyCredentials);
|
||||
},
|
||||
};
|
||||
|
||||
// TODO: We can request population of the user info if absent or expired.
|
||||
|
||||
let user_info = match mj_creds_manager.maybe_copy_user_info() {
|
||||
Ok(Some(user_info)) => user_info,
|
||||
Ok(None) => {
|
||||
return Err(InternalImageError::NeedsMidjourneyUserInfo);
|
||||
}
|
||||
Err(err) => {
|
||||
error!("Error reading Midjourney user info: {:?}", err);
|
||||
return Err(InternalImageError::NeedsMidjourneyUserInfo);
|
||||
},
|
||||
};
|
||||
|
||||
let channel_id = match user_info.user_id {
|
||||
Some(user_id) => ChannelId::UserId(user_id),
|
||||
None => {
|
||||
error!("Midjourney user info does not contain a user ID.");
|
||||
return Err(InternalImageError::NeedsMidjourneyUserId);
|
||||
}
|
||||
};
|
||||
|
||||
info!("Calling midjourney ...");
|
||||
|
||||
let cookie_header = creds.to_cookie_string();
|
||||
|
||||
let prompt = request.prompt
|
||||
.as_deref()
|
||||
.unwrap_or("");
|
||||
|
||||
let result = text_to_image(TextToImageRequest {
|
||||
prompt,
|
||||
channel_id: &channel_id,
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
cookie_header,
|
||||
}).await;
|
||||
|
||||
let result = match result {
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
error!("Failed to use MidJourney: {:?}", err);
|
||||
return Err(InternalImageError::MidjourneyError(err));
|
||||
}
|
||||
};
|
||||
|
||||
let job_id = match result.maybe_job_id {
|
||||
Some(job_id) => job_id,
|
||||
None => {
|
||||
error!("Failed to enqueue MidJourney: No job ID returned.");
|
||||
return handle_midjourney_errors(app, result.maybe_errors);
|
||||
}
|
||||
};
|
||||
|
||||
info!("Successfully enqueued MidJourney. Job token: {}", job_id);
|
||||
|
||||
Ok(TaskEnqueueSuccess {
|
||||
provider: GenerationProvider::Midjourney,
|
||||
model: Some(GenerationModel::Midjourney),
|
||||
provider_job_id: Some(job_id),
|
||||
task_type: TaskType::ImageGeneration,
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_midjourney_errors(
|
||||
app: &AppHandle,
|
||||
maybe_errors: Option<Vec<TextToImageError>>
|
||||
) -> Result<TaskEnqueueSuccess, InternalImageError> {
|
||||
if let Some(errors) = maybe_errors {
|
||||
if !errors.is_empty() {
|
||||
let messages: Vec<String> = errors.iter()
|
||||
.map(|e| format!("{:?}", e))
|
||||
.collect();
|
||||
|
||||
let combined_message = messages.join("; ");
|
||||
|
||||
let event = FlashUserInputErrorEvent {
|
||||
message: format!("Midjourney Error: {}", combined_message),
|
||||
};
|
||||
|
||||
if let Err(err) = event.send(&app) {
|
||||
error!("Failed to send FlashUserInputErrorEvent: {:?}", err); // Fail open
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(InternalImageError::MidjourneyJobEnqueueFailed)
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub (super) mod handle_midjourney;
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod enqueue_text_to_image_command;
|
||||
pub mod generic;
|
||||
pub mod internal_image_error;
|
||||
mod gpt_image_1;
|
||||
mod gpt_image_1;
|
||||
mod midjourney;
|
||||
+1
@@ -129,6 +129,7 @@ async fn upload_image_from_base64_bytes(
|
||||
maybe_creds: Some(&creds),
|
||||
path: temp_file,
|
||||
is_intermediate_system_file: true, // NB: Probably not essential to keep this.
|
||||
maybe_prompt_token: None, // NB: Not used for bg removal.
|
||||
}).await
|
||||
.map_err(|err| {
|
||||
error!("Failed to upload image media file: {:?}", err);
|
||||
|
||||
@@ -31,6 +31,7 @@ impl TaskEnqueueSuccess{
|
||||
match self.provider {
|
||||
GenerationProvider::Artcraft => GenerationServiceProvider::Artcraft,
|
||||
GenerationProvider::Fal => GenerationServiceProvider::Fal,
|
||||
GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney,
|
||||
GenerationProvider::Sora => GenerationServiceProvider::Sora,
|
||||
}
|
||||
}
|
||||
@@ -69,6 +70,7 @@ impl TaskEnqueueSuccess{
|
||||
Some(GenerationModel::Veo2) => Some(TaskModelType::Veo2),
|
||||
Some(GenerationModel::Hunyuan3d2_0) => Some(TaskModelType::Hunyuan3d2_0),
|
||||
Some(GenerationModel::Hunyuan3d2_1) => Some(TaskModelType::Hunyuan3d2_1),
|
||||
Some(GenerationModel::Midjourney) => Some(TaskModelType::Midjourney), // NB: This is a generic Midjourney model, version unknown.
|
||||
|
||||
// TODO: These seem wrong -
|
||||
Some(GenerationModel::Kling1_6) => Some(TaskModelType::Kling16Pro), // NB: `VideoModel::Kling16Pro`.
|
||||
|
||||
@@ -1 +1,2 @@
|
||||
pub mod canvas_background_removal_complete_event;
|
||||
pub mod canvas_background_removal_complete_event;
|
||||
pub mod refresh_account_state_event;
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::tauri::ux::tauri_event_name::TauriEventName;
|
||||
use serde_derive::Serialize;
|
||||
|
||||
/// Send a signal to the frontend that an account state has changed: login,
|
||||
/// logout, invalidated, etc.
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub struct RefreshAccountStateEvent {
|
||||
/// The frontend can optionally decide whether it wants to refresh based
|
||||
/// on the impacted provider account.
|
||||
pub provider: Option<GenerationProvider>,
|
||||
}
|
||||
|
||||
impl BasicSendableEvent for RefreshAccountStateEvent {
|
||||
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::RefreshAccountStateEvent;
|
||||
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success;
|
||||
}
|
||||
@@ -5,6 +5,7 @@ use serde::Serialize;
|
||||
pub enum GenerationServiceProvider {
|
||||
Artcraft,
|
||||
Fal,
|
||||
Midjourney,
|
||||
Sora,
|
||||
}
|
||||
|
||||
@@ -30,6 +31,10 @@ pub enum GenerationModel {
|
||||
#[serde(rename = "hunyuan_3d_2_1")]
|
||||
Hunyuan3d2_1,
|
||||
|
||||
// Generic Midjourney model, version unknown.
|
||||
#[serde(rename = "midjourney")]
|
||||
Midjourney,
|
||||
|
||||
// TODO: Should be Kling16Pro
|
||||
#[serde(rename = "kling_1.6")]
|
||||
Kling1_6,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub mod basic_sendable_event_trait;
|
||||
pub mod functional_events;
|
||||
pub mod generation_events;
|
||||
pub mod sendable_event_error;
|
||||
pub mod sendable_event_trait;
|
||||
pub mod functional_events;
|
||||
pub mod warning_events;
|
||||
+14
@@ -0,0 +1,14 @@
|
||||
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
|
||||
use enums::tauri::ux::tauri_event_name::TauriEventName;
|
||||
use serde_derive::Serialize;
|
||||
|
||||
#[derive(Clone, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub struct FlashUserInputErrorEvent {
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
impl BasicSendableEvent for FlashUserInputErrorEvent {
|
||||
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::FlashUserInputErrorEvent;
|
||||
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Failure;
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod flash_user_input_error_event;
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::core::lifecycle::startup::tasks::bootstrap_task_database::bootstrap_task_database;
|
||||
use crate::core::lifecycle::startup::tasks::initially_size_and_position_windows::initially_size_and_position_windows;
|
||||
use crate::core::lifecycle::startup::tasks::load_provider_priority_state::load_provider_priority_state;
|
||||
use crate::core::lifecycle::startup::tasks::set_app_log_level::set_app_log_level;
|
||||
use crate::core::lifecycle::startup::tasks::initially_size_and_position_windows::initially_size_and_position_windows;
|
||||
use crate::core::lifecycle::startup::tasks::spawn_discord_presence_thread::spawn_discord_presence_thread;
|
||||
use crate::core::lifecycle::startup::tasks::spawn_fal_task_polling_thread::spawn_fal_task_polling_thread;
|
||||
use crate::core::lifecycle::startup::tasks::spawn_main_window_thread::spawn_main_window_thread;
|
||||
@@ -11,6 +11,8 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
|
||||
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread;
|
||||
use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
@@ -26,6 +28,7 @@ pub async fn handle_tauri_startup(
|
||||
sora_task_queue: SoraTaskQueue,
|
||||
fal_credential_manager: FalCredentialManager,
|
||||
fal_task_queue: FalTaskQueue,
|
||||
mj_creds_manager: MidjourneyCredentialManager,
|
||||
) -> AnyhowResult<()> {
|
||||
|
||||
set_app_log_level(
|
||||
@@ -72,6 +75,15 @@ pub async fn handle_tauri_startup(
|
||||
&storyteller_creds_manager,
|
||||
)?;
|
||||
|
||||
tauri::async_runtime::spawn(midjourney_long_polling_thread(
|
||||
app.clone(),
|
||||
app_env_configs.clone(),
|
||||
root.clone(),
|
||||
task_database.clone(),
|
||||
mj_creds_manager.clone(),
|
||||
storyteller_creds_manager.clone(),
|
||||
));
|
||||
|
||||
spawn_discord_presence_thread()?;
|
||||
|
||||
initially_size_and_position_windows(&app, &root);
|
||||
|
||||
@@ -17,4 +17,8 @@ pub enum ImageModel {
|
||||
GptImage1,
|
||||
#[serde(rename = "recraft_3")]
|
||||
Recraft3,
|
||||
|
||||
// Generic Midjourney model, version unknown.
|
||||
#[serde(rename = "midjourney")]
|
||||
Midjourney,
|
||||
}
|
||||
|
||||
@@ -44,4 +44,8 @@ impl AppCredentialsDir {
|
||||
pub fn get_fal_api_key_file_path(&self) -> PathBuf {
|
||||
self.path.join("fal_api_key.txt")
|
||||
}
|
||||
|
||||
pub fn get_midjourney_state_path(&self) -> PathBuf {
|
||||
self.path.join("midjourney_state.json")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,7 +22,9 @@ impl DataSubdir for AppStateDir{
|
||||
|
||||
impl AppStateDir {
|
||||
pub fn get_tasks_sqlite_database_path(&self) -> PathBuf {
|
||||
self.path.join("tasks_v2.sqlite")
|
||||
// NB: This must be incremented on any change to the database schema (including comments).
|
||||
// The database houses ephemeral content and will be migrated automatically.
|
||||
self.path.join("tasks_v3.sqlite")
|
||||
}
|
||||
|
||||
pub fn get_window_size_config_file(&self) -> PathBuf {
|
||||
|
||||
@@ -6,6 +6,7 @@ pub fn to_generation_service_provider(provider: GenerationProvider) -> Generatio
|
||||
match provider {
|
||||
GenerationProvider::Artcraft => GenerationServiceProvider::Artcraft,
|
||||
GenerationProvider::Fal => GenerationServiceProvider::Fal,
|
||||
GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney,
|
||||
GenerationProvider::Sora => GenerationServiceProvider::Sora,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
use anyhow::anyhow;
|
||||
use tauri::WebviewWindow;
|
||||
use errors::AnyhowResult;
|
||||
|
||||
pub fn get_webview_window_hostname(webview: &WebviewWindow) -> AnyhowResult<String> {
|
||||
let url = webview.url()?;
|
||||
let url_hostname= url.host()
|
||||
.ok_or(anyhow!("no host in url"))?
|
||||
.to_string();
|
||||
Ok(url_hostname)
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod get_webview_window_hostname;
|
||||
pub mod get_window_size_heuristic;
|
||||
pub mod position_main_window;
|
||||
pub mod resize_main_window;
|
||||
|
||||
@@ -7,8 +7,8 @@ use crate::core::commands::app_preferences::get_app_preferences_command::get_app
|
||||
use crate::core::commands::app_preferences::update_app_preference_command::update_app_preferences_command;
|
||||
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::enqueue_text_to_image_command;
|
||||
use crate::core::commands::enqueue::image_bg_removal::enqueue_image_bg_removal_command::enqueue_image_bg_removal_command;
|
||||
use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::enqueue_image_inpaint_command;
|
||||
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::enqueue_contextual_edit_image_command;
|
||||
use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::enqueue_image_inpaint_command;
|
||||
use crate::core::commands::enqueue::object::enqueue_image_to_3d_object_command::enqueue_image_to_3d_object_command;
|
||||
use crate::core::commands::enqueue::video::enqueue_image_to_video_command::enqueue_image_to_video_command;
|
||||
use crate::core::commands::flip_image::flip_image;
|
||||
@@ -32,6 +32,10 @@ use crate::services::fal::commands::set_fal_api_key_command::set_fal_api_key_com
|
||||
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
|
||||
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
|
||||
use crate::services::fal::threads::fal_task_polling_thread::fal_task_polling_thread;
|
||||
use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command;
|
||||
use crate::services::midjourney::commands::midjourney_get_credential_info_command::midjourney_get_credential_info_command;
|
||||
use crate::services::midjourney::commands::midjourney_open_login_command::midjourney_open_login_command;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::sora::commands::check_sora_session_command::check_sora_session_command;
|
||||
use crate::services::sora::commands::open_sora_login_command::open_sora_login_command;
|
||||
use crate::services::sora::commands::sora_image_generation_command::sora_image_generation_command;
|
||||
@@ -85,7 +89,10 @@ pub fn run() {
|
||||
.expect("AppEnvConfigs should be loaded from disk");
|
||||
|
||||
let app_env_configs_2 = app_env_configs.clone();
|
||||
|
||||
|
||||
let midjourney_creds_manager = MidjourneyCredentialManager::initialize_from_disk_infallible(&app_data_root);
|
||||
let midjourney_creds_manager_2 = midjourney_creds_manager.clone();
|
||||
|
||||
println!("Initializing backend runtime...");
|
||||
|
||||
tauri::Builder::default()
|
||||
@@ -123,6 +130,7 @@ pub fn run() {
|
||||
sora_tasks,
|
||||
fal_creds,
|
||||
fal_tasks,
|
||||
midjourney_creds_manager_2,
|
||||
).await;
|
||||
|
||||
if let Err(err) = result {
|
||||
@@ -138,6 +146,7 @@ pub fn run() {
|
||||
.manage(app_preferences)
|
||||
.manage(fal_creds_manager)
|
||||
.manage(fal_task_queue)
|
||||
.manage(midjourney_creds_manager)
|
||||
.manage(sora_creds_manager)
|
||||
.manage(sora_task_queue)
|
||||
.manage(storyteller_creds_manager_3)
|
||||
@@ -159,6 +168,9 @@ pub fn run() {
|
||||
get_provider_order_command,
|
||||
load_without_cors_command,
|
||||
open_sora_login_command,
|
||||
midjourney_clear_credentials_command,
|
||||
midjourney_get_credential_info_command,
|
||||
midjourney_open_login_command,
|
||||
platform_info_command,
|
||||
set_fal_api_key_command,
|
||||
set_provider_order_command,
|
||||
|
||||
@@ -261,6 +261,7 @@ pub async fn remove_background(
|
||||
maybe_creds: Some(&creds),
|
||||
path: result_filename,
|
||||
is_intermediate_system_file: true, // NB: Probably not essential to keep this.
|
||||
maybe_prompt_token: None, // NB: Not used for bg removal.
|
||||
}).await?;
|
||||
|
||||
// TODO: Don't re-request to simply build MediaLinks (or CDN URL). Get those from the upload API in one turn.
|
||||
|
||||
@@ -125,7 +125,8 @@ async fn polling_loop(
|
||||
let result = upload_media_file_from_file(
|
||||
&app_env_configs.storyteller_host,
|
||||
Some(&storyteller_creds),
|
||||
download_file
|
||||
download_file,
|
||||
None, // TODO: Should probably save prompts.
|
||||
).await;
|
||||
|
||||
match result {
|
||||
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
use crate::core::commands::response::shorthand::{Response, SimpleResponse};
|
||||
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use errors::AnyhowResult;
|
||||
use fal_client::creds::fal_api_key::FalApiKey;
|
||||
use log::{error, info};
|
||||
use serde_derive::Deserialize;
|
||||
use tauri::State;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn midjourney_clear_credentials_command(
|
||||
root: State<'_, AppDataRoot>,
|
||||
creds_manager: State<'_, MidjourneyCredentialManager>,
|
||||
) -> SimpleResponse {
|
||||
info!("midjourney_clear_credentials_command called");
|
||||
|
||||
clear_creds(&root, &creds_manager)
|
||||
.map_err(|err| {
|
||||
error!("Error clearing creds: {:?}", err);
|
||||
"error clearing creds"
|
||||
})?;
|
||||
|
||||
Ok(().into())
|
||||
}
|
||||
|
||||
fn clear_creds(
|
||||
root: &AppDataRoot,
|
||||
creds: &MidjourneyCredentialManager,
|
||||
) -> AnyhowResult<()> {
|
||||
|
||||
creds.clear_credentials()?;
|
||||
creds.persist_to_disk()?;
|
||||
|
||||
let creds_path = root.credentials_dir().get_midjourney_state_path();
|
||||
if creds_path.exists() {
|
||||
std::fs::remove_file(creds_path)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
use crate::core::commands::response::shorthand::{Response, ResponseOrErrorMessage};
|
||||
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
|
||||
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use errors::AnyhowResult;
|
||||
use log::{error, info};
|
||||
use serde_derive::Serialize;
|
||||
use tauri::State;
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct MidjourneyGetCredentialInfoResponse {
|
||||
pub maybe_email: Option<String>,
|
||||
pub can_clear_state: bool,
|
||||
}
|
||||
|
||||
impl SerializeMarker for MidjourneyGetCredentialInfoResponse {}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn midjourney_get_credential_info_command(
|
||||
creds_manager: State<'_, MidjourneyCredentialManager>,
|
||||
) -> ResponseOrErrorMessage<MidjourneyGetCredentialInfoResponse> {
|
||||
info!("midjourney_get_credential_info_command called");
|
||||
|
||||
let info = get_info(&creds_manager)
|
||||
.map_err(|err| {
|
||||
error!("Error getting info: {:?}", err);
|
||||
"error getting info"
|
||||
})?;
|
||||
|
||||
Ok(info.into())
|
||||
}
|
||||
|
||||
fn get_info(
|
||||
creds: &MidjourneyCredentialManager,
|
||||
) -> AnyhowResult<MidjourneyGetCredentialInfoResponse> {
|
||||
let mut can_clear_state = true;
|
||||
|
||||
let maybe_cookies = creds.maybe_copy_cookie_store()?;
|
||||
let maybe_user_info = creds.maybe_copy_user_info()?;
|
||||
|
||||
if maybe_cookies.is_none() && maybe_user_info.is_none() {
|
||||
can_clear_state = false;
|
||||
}
|
||||
|
||||
let maybe_email = maybe_user_info
|
||||
.map(|info| info.email.clone())
|
||||
.flatten();
|
||||
|
||||
Ok(MidjourneyGetCredentialInfoResponse {
|
||||
maybe_email,
|
||||
can_clear_state,
|
||||
})
|
||||
}
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::midjourney::windows::open_midjourney_login_window::open_midjourney_login_window;
|
||||
use crate::services::sora::windows::sora_login_window::open_sora_login_window::open_sora_login_window;
|
||||
use errors::AnyhowResult;
|
||||
use log::{error, info};
|
||||
use tauri::{AppHandle, State};
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn midjourney_open_login_command(
|
||||
app: AppHandle,
|
||||
app_data_root: State<'_, AppDataRoot>,
|
||||
mj_creds_manager: State<'_, MidjourneyCredentialManager>,
|
||||
) -> Result<String, String> {
|
||||
info!("midjourney_open_login_command called");
|
||||
|
||||
do_open_login(&app, &app_data_root, &mj_creds_manager)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
error!("Error opening login: {:?}", err);
|
||||
format!("Error opening login: {:?}", err)
|
||||
})?;
|
||||
|
||||
Ok("result".to_string())
|
||||
}
|
||||
|
||||
async fn do_open_login(
|
||||
app: &AppHandle,
|
||||
app_data_root: &AppDataRoot,
|
||||
mj_creds_manager: &MidjourneyCredentialManager,
|
||||
) -> AnyhowResult<()> {
|
||||
info!("Building login window...");
|
||||
|
||||
open_midjourney_login_window(app, app_data_root, mj_creds_manager).await?;
|
||||
|
||||
info!("Done.");
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod midjourney_clear_credentials_command;
|
||||
pub mod midjourney_get_credential_info_command;
|
||||
pub mod midjourney_open_login_command;
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod commands;
|
||||
pub mod state;
|
||||
pub mod threads;
|
||||
pub mod utils;
|
||||
pub mod windows;
|
||||
+196
@@ -0,0 +1,196 @@
|
||||
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, SERIALIZABLE_MIDJOURNEY_STATE_VERSION};
|
||||
use crate::services::storyteller::state::read_storyteller_credentials_from_disk::read_storyteller_credentials_from_disk;
|
||||
use crate::services::storyteller::state::storyteller_credential_holder::StorytellerCredentialHolder;
|
||||
use cookie_store::cookie_store::CookieStore;
|
||||
use errors::AnyhowResult;
|
||||
use log::{info, warn};
|
||||
use std::fs::read_to_string;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
|
||||
use storyteller_client::credentials::storyteller_avt_cookie::StorytellerAvtCookie;
|
||||
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use storyteller_client::credentials::storyteller_session_cookie::StorytellerSessionCookie;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MidjourneyCredentialManager {
|
||||
// TODO: Put last write/last read timestamps on these.
|
||||
cookies: Arc<RwLock<Option<CookieStore>>>,
|
||||
user_info: Arc<RwLock<Option<MidjourneyUserInfo>>>,
|
||||
app_data_root: AppDataRoot,
|
||||
}
|
||||
|
||||
impl MidjourneyCredentialManager {
|
||||
pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self {
|
||||
Self {
|
||||
cookies: Arc::new(RwLock::new(None)),
|
||||
user_info: Arc::new(RwLock::new(None)),
|
||||
app_data_root: app_data_root.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self {
|
||||
let mut cookies;
|
||||
let mut user_info;
|
||||
|
||||
match read_midjourney_state_from_disk(app_data_root) {
|
||||
Err(err) => {
|
||||
warn!("Failed to read midjourney state from disk: {:?}", err);
|
||||
cookies = Arc::new(RwLock::new(None));
|
||||
user_info = Arc::new(RwLock::new(None));
|
||||
},
|
||||
Ok(None) => {
|
||||
cookies = Arc::new(RwLock::new(None));
|
||||
user_info = Arc::new(RwLock::new(None));
|
||||
}
|
||||
Ok(Some(state)) => {
|
||||
let maybe_cookies = state.user_cookies
|
||||
.map(|cookies| cookies.to_cookie_store());
|
||||
let maybe_user_info = state.user_info
|
||||
.map(|info| info.to_user_info());
|
||||
cookies = Arc::new(RwLock::new(maybe_cookies));
|
||||
user_info = Arc::new(RwLock::new(maybe_user_info));
|
||||
}
|
||||
};
|
||||
|
||||
Self {
|
||||
cookies,
|
||||
user_info,
|
||||
app_data_root: app_data_root.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
|
||||
match self.cookies.read() {
|
||||
Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
|
||||
Ok(store) => Ok(store.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn maybe_copy_user_info(&self) -> anyhow::Result<Option<MidjourneyUserInfo>> {
|
||||
match self.user_info.read() {
|
||||
Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
|
||||
Ok(info) => Ok(info.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> {
|
||||
match self.cookies.write() {
|
||||
Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
|
||||
Ok(mut current_store) => {
|
||||
*current_store = Some(store);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn replace_user_info(&self, user_info: MidjourneyUserInfo) -> anyhow::Result<()> {
|
||||
match self.user_info.write() {
|
||||
Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
|
||||
Ok(mut current_info) => {
|
||||
*current_info = Some(user_info);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// NB: This is just a heuristic. We'll add better checks later.
|
||||
pub fn session_appears_active(&self) -> anyhow::Result<bool> {
|
||||
let maybe_cookies = match self.cookies.read() {
|
||||
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
|
||||
Ok(store) => store.clone(),
|
||||
};
|
||||
|
||||
let maybe_user_info = match self.user_info.read() {
|
||||
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
|
||||
Ok(info) => info.clone(),
|
||||
};
|
||||
|
||||
if maybe_cookies.is_none() || maybe_user_info.is_none() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let cookies = maybe_cookies.unwrap();
|
||||
let user_info = maybe_user_info.unwrap();
|
||||
|
||||
// NB: We consider the session active if we have auth cookies and a user id.
|
||||
let has_user_id = user_info.user_id.is_some();
|
||||
let maybe_has_auth_cookies = cookie_store_has_auth_cookies(&cookies);
|
||||
|
||||
// Misc cookies without login cookies are ~1055 length
|
||||
// AUTH_I is ~1500 length
|
||||
// AUTH_R is ~500 length
|
||||
let maybe_has_big_enough_cookie = cookies.calculate_approx_cookie_character_length() > 2100;
|
||||
|
||||
// TODO: Heuristic should count.
|
||||
// TODO: Consolidate with "login window thread" logic.
|
||||
// TODO: Check timestamp of last valid request.
|
||||
Ok(has_user_id && maybe_has_auth_cookies && maybe_has_big_enough_cookie && maybe_has_big_enough_cookie)
|
||||
}
|
||||
|
||||
pub fn clear_credentials(&self) -> anyhow::Result<()> {
|
||||
match self.cookies.write() {
|
||||
Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
|
||||
Ok(mut store) => *store = None,
|
||||
}
|
||||
match self.user_info.write() {
|
||||
Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
|
||||
Ok(mut info) => *info = None,
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn persist_to_disk(&self) -> anyhow::Result<()> {
|
||||
let maybe_cookies = match self.cookies.read() {
|
||||
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
|
||||
Ok(store) => store.clone(),
|
||||
};
|
||||
|
||||
let maybe_user_info = match self.user_info.read() {
|
||||
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
|
||||
Ok(info) => info.clone(),
|
||||
};
|
||||
|
||||
let maybe_cookies= maybe_cookies
|
||||
.map(|cookies| cookies.to_serializable());
|
||||
|
||||
let maybe_user_info= maybe_user_info
|
||||
.map(|info| info.to_serializable());
|
||||
|
||||
if maybe_cookies.is_none() && maybe_user_info.is_none() {
|
||||
info!("Nothing to write to disk, skipping.");
|
||||
return Ok(());
|
||||
|
||||
}
|
||||
|
||||
let state = SerializableMidjourneyState {
|
||||
version: SERIALIZABLE_MIDJOURNEY_STATE_VERSION,
|
||||
user_cookies: maybe_cookies,
|
||||
user_info: maybe_user_info,
|
||||
};
|
||||
|
||||
let path = self.app_data_root.credentials_dir().get_midjourney_state_path();
|
||||
let serialized = serde_json::to_string(&state)?;
|
||||
|
||||
std::fs::write(path, serialized)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn read_midjourney_state_from_disk(root: &AppDataRoot) -> AnyhowResult<Option<SerializableMidjourneyState>> {
|
||||
let midjourney_state_path= root.credentials_dir().get_midjourney_state_path();
|
||||
|
||||
if !midjourney_state_path.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let contents = read_to_string(&midjourney_state_path)?
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
let state: SerializableMidjourneyState = serde_json::from_str(&contents)?;
|
||||
|
||||
Ok(Some(state))
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use midjourney_client::recipes::get_user_info::GetUserInfoResponse;
|
||||
use crate::services::midjourney::state::serializable_midjourney_state::SerializableMidjourneyUserInfo;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MidjourneyUserInfo {
|
||||
pub user_id: Option<MidjourneyUserId>,
|
||||
pub email: Option<String>,
|
||||
pub websocket_token: Option<String>,
|
||||
}
|
||||
|
||||
impl MidjourneyUserInfo {
|
||||
pub fn to_serializable(&self) -> SerializableMidjourneyUserInfo {
|
||||
SerializableMidjourneyUserInfo {
|
||||
user_id: self.user_id
|
||||
.as_ref()
|
||||
.map(|id| id.to_string()),
|
||||
email: self.email.clone(),
|
||||
websocket_token: self.websocket_token.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_api_response(response: GetUserInfoResponse) -> Self {
|
||||
Self {
|
||||
user_id: response.user_id,
|
||||
email: response.email,
|
||||
websocket_token: response.websocket_token,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
pub (super) mod serializable_midjourney_state;
|
||||
pub mod midjourney_credential_manager;
|
||||
pub mod midjourney_user_info;
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
|
||||
use cookie_store::serialized_cookie_store::SerializableCookieStore;
|
||||
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use serde_derive::{Deserialize, Serialize};
|
||||
|
||||
pub (super) const SERIALIZABLE_MIDJOURNEY_STATE_VERSION: u32 = 1;
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub (super) struct SerializableMidjourneyState {
|
||||
pub (super) version: u32,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) user_cookies: Option<SerializableCookieStore>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) user_info: Option<SerializableMidjourneyUserInfo>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub (super) struct SerializableMidjourneyUserInfo {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) user_id: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) email: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) websocket_token: Option<String>,
|
||||
}
|
||||
|
||||
impl SerializableMidjourneyUserInfo {
|
||||
pub fn to_user_info(&self) -> MidjourneyUserInfo {
|
||||
MidjourneyUserInfo {
|
||||
user_id: self.user_id
|
||||
.as_deref()
|
||||
.map(|id| MidjourneyUserId::from_str(id)),
|
||||
email: self.email.clone(),
|
||||
websocket_token: self.websocket_token.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
+323
@@ -0,0 +1,323 @@
|
||||
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
|
||||
use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
|
||||
use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
|
||||
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
use crate::core::state::task_database::TaskDatabase;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image;
|
||||
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest;
|
||||
use cookie_store::cookie_store::CookieStore;
|
||||
use enums::by_table::prompts::prompt_type::PromptType;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use enums::common::model_type::ModelType;
|
||||
use enums::tauri::tasks::task_status::TaskStatus;
|
||||
use errors::AnyhowResult;
|
||||
use idempotency::uuid::generate_random_uuid;
|
||||
use log::{error, info};
|
||||
use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
|
||||
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
|
||||
use midjourney_client::endpoints::imagine::{imagine, ImagineRequest, MidjourneyJobType};
|
||||
use midjourney_client::utils::get_image_url::get_image_url;
|
||||
use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
|
||||
use once_cell::sync::Lazy;
|
||||
use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, Task, TaskList};
|
||||
use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fs::File;
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
|
||||
use storyteller_client::error::api_error::ApiError;
|
||||
use storyteller_client::error::storyteller_error::StorytellerError;
|
||||
use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
|
||||
use storyteller_client::prompts::create_prompt::create_prompt;
|
||||
use tauri::AppHandle;
|
||||
use url::Url;
|
||||
|
||||
static PENDING_STATUSES : Lazy<HashSet<TaskStatus>> = Lazy::new(|| {
|
||||
let mut statuses = HashSet::new();
|
||||
statuses.insert(TaskStatus::Pending);
|
||||
statuses.insert(TaskStatus::Started);
|
||||
statuses.insert(TaskStatus::AttemptFailed);
|
||||
statuses
|
||||
});
|
||||
|
||||
/// This thread is responsible for picking up tasks that fell through the cracks of
|
||||
/// the faster websocket thread.
|
||||
pub async fn midjourney_long_polling_thread(
|
||||
app_handle: AppHandle,
|
||||
app_env_configs: AppEnvConfigs,
|
||||
app_data_root: AppDataRoot,
|
||||
task_database: TaskDatabase,
|
||||
creds: MidjourneyCredentialManager,
|
||||
storyteller_creds_manager: StorytellerCredentialManager,
|
||||
) -> ! {
|
||||
loop {
|
||||
let res = polling_loop(
|
||||
&app_handle,
|
||||
&app_env_configs,
|
||||
&app_data_root,
|
||||
&task_database,
|
||||
&creds,
|
||||
&storyteller_creds_manager,
|
||||
).await;
|
||||
if let Err(err) = res {
|
||||
error!("An error occurred: {:?}", err);
|
||||
}
|
||||
// NB: Only sleep if an error occurs.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn polling_loop(
|
||||
app_handle: &AppHandle,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
app_data_root: &AppDataRoot,
|
||||
task_database: &TaskDatabase,
|
||||
creds: &MidjourneyCredentialManager,
|
||||
storyteller_creds_manager: &StorytellerCredentialManager,
|
||||
) -> AnyhowResult<()> {
|
||||
loop {
|
||||
if !creds.session_appears_active()? {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// TODO: Graceful wait, fix this long function body
|
||||
let storyteller_creds = match storyteller_creds_manager.get_credentials()? {
|
||||
Some(creds) => creds,
|
||||
None => {
|
||||
error!("No Storyteller credentials found. Cannot proceed with Midjourney polling.");
|
||||
tokio::time::sleep(std::time::Duration::from_millis(5_000)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let cookies = creds.maybe_copy_cookie_store()?;
|
||||
|
||||
let cookies = match cookies {
|
||||
Some(cookies) => cookies,
|
||||
None => {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let user_info = creds.maybe_copy_user_info()?;
|
||||
|
||||
let maybe_user_id = user_info
|
||||
.map(|info| info.user_id)
|
||||
.flatten();
|
||||
|
||||
let user_id = match maybe_user_id {
|
||||
Some(user_id) => user_id,
|
||||
None => {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs {
|
||||
db: task_database.get_connection(),
|
||||
provider: GenerationProvider::Midjourney,
|
||||
task_statuses: &PENDING_STATUSES,
|
||||
}).await?;
|
||||
|
||||
check_midjourney_tasks(
|
||||
app_handle,
|
||||
app_env_configs,
|
||||
app_data_root,
|
||||
task_database,
|
||||
&cookies,
|
||||
&user_id,
|
||||
&storyteller_creds,
|
||||
local_tasks,
|
||||
).await?;
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn check_midjourney_tasks(
|
||||
app_handle: &AppHandle,
|
||||
app_env_configs: &AppEnvConfigs,
|
||||
app_data_root: &AppDataRoot,
|
||||
task_database: &TaskDatabase,
|
||||
mj_cookies: &CookieStore,
|
||||
mj_user_id: &MidjourneyUserId,
|
||||
storyteller_creds: &StorytellerCredentialSet,
|
||||
local_tasks: TaskList,
|
||||
) -> AnyhowResult<()> {
|
||||
let local_tasks = local_tasks.tasks;
|
||||
|
||||
if local_tasks.is_empty() {
|
||||
return Ok(())
|
||||
}
|
||||
|
||||
// Map of Midjourney Job ID to Local Task.
|
||||
let local_tasks_by_midjourney_job_id = local_tasks.iter()
|
||||
.filter_map(|task| {
|
||||
if let Some(provider_job_id) = &task.provider_job_id {
|
||||
Some((provider_job_id.clone(), task.clone()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect::<HashMap<String, Task>>();
|
||||
|
||||
let cookie_header = mj_cookies.to_cookie_string();
|
||||
|
||||
let midjourney_result = imagine(ImagineRequest {
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
cookie_header,
|
||||
user_id: mj_user_id,
|
||||
page_size: None,
|
||||
}).await?;
|
||||
|
||||
let midjourney_items = midjourney_result.items;
|
||||
|
||||
let midjourney_items_by_id = {
|
||||
let mut hash = HashMap::new();
|
||||
for item in midjourney_items.iter() {
|
||||
if let Some(id) = &item.id {
|
||||
hash.insert(id.to_string(), item.clone());
|
||||
}
|
||||
}
|
||||
hash
|
||||
};
|
||||
|
||||
// TODO: If we introduce another job polling mechanism, we may need to handle concurrency.
|
||||
// One idea might be to add a new job state that acts as an optimistic lock
|
||||
|
||||
let image_downloader = ImageDownloaderClient::create()?;
|
||||
|
||||
for (midjourney_job_id, local_task) in local_tasks_by_midjourney_job_id.iter() {
|
||||
// TODO: Copy prompt from this.
|
||||
let midjourney_item = match midjourney_items_by_id.get(midjourney_job_id) {
|
||||
Some(item) => item,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let model_type = match midjourney_item.job_type {
|
||||
Some(MidjourneyJobType::V6Diffusion) => ModelType::MidjourneyV6,
|
||||
Some(MidjourneyJobType::V6p1Diffusion) => ModelType::MidjourneyV6p1,
|
||||
Some(MidjourneyJobType::V6p1RawDiffusion) => ModelType::MidjourneyV6p1Raw,
|
||||
Some(MidjourneyJobType::V7Diffusion) => ModelType::MidjourneyV7,
|
||||
Some(MidjourneyJobType::V7RawDiffusion) => ModelType::MidjourneyV7Raw,
|
||||
Some(MidjourneyJobType::V7DraftDiffusion) => ModelType::MidjourneyV7Draft,
|
||||
Some(MidjourneyJobType::V7DraftRawDiffusion) => ModelType::MidjourneyV7DraftRaw,
|
||||
Some(MidjourneyJobType::Other(ref other)) => {
|
||||
info!("Unknown Midjourney job type (for job id {}): {}", midjourney_job_id, other);
|
||||
ModelType::Midjourney
|
||||
},
|
||||
_ => ModelType::Midjourney,
|
||||
};
|
||||
|
||||
let request = CreatePromptRequest {
|
||||
uuid_idempotency_token: generate_random_uuid(),
|
||||
positive_prompt: midjourney_item.full_command.clone(),
|
||||
negative_prompt: None,
|
||||
model_type: Some(model_type),
|
||||
generation_provider: Some(GenerationProvider::Midjourney),
|
||||
};
|
||||
|
||||
let prompt_response = create_prompt(
|
||||
&app_env_configs.storyteller_host,
|
||||
Some(storyteller_creds),
|
||||
request
|
||||
).await?;
|
||||
|
||||
info!("Created prompt: {:?}", &prompt_response.prompt_token);
|
||||
|
||||
for index in 0..4 {
|
||||
info!("Downloading generated Midjourney file...");
|
||||
|
||||
let download_path = download_midjourney_image(
|
||||
&image_downloader,
|
||||
midjourney_job_id,
|
||||
index,
|
||||
app_data_root
|
||||
).await?;
|
||||
|
||||
let mut wait_delay = 0;
|
||||
|
||||
loop {
|
||||
info!("Uploading to backend...");
|
||||
|
||||
// TODO: media_files.origin_category
|
||||
// TODO: media_files.maybe_prompt_token
|
||||
// TODO: media_files.maybe_generation_provider
|
||||
// TODO: media_files.maybe_origin_model_type
|
||||
// TODO: media_files.maybe_origin_model_token (sref?)
|
||||
// TODO: media_files.is_batch_generated
|
||||
// TODO: media_files.maybe_batch_token
|
||||
// TODO: media_files.is_user_upload
|
||||
|
||||
// TODO: batch_generations.token
|
||||
// TODO: batch_generations.entity_type
|
||||
// TODO: batch_generations.entity_token
|
||||
|
||||
let result = upload_image_media_file_from_file(UploadImageFromFileArgs {
|
||||
api_host: &app_env_configs.storyteller_host,
|
||||
maybe_creds: Some(&storyteller_creds),
|
||||
path: &download_path,
|
||||
is_intermediate_system_file: false,
|
||||
maybe_prompt_token: Some(&prompt_response.prompt_token),
|
||||
}).await;
|
||||
|
||||
match result {
|
||||
Ok(media_file) => {
|
||||
info!("Successfully uploaded to backend: {:?}", media_file);
|
||||
break;
|
||||
},
|
||||
Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => {
|
||||
error!("Too many requests, retrying upload after delay...");
|
||||
// If we hit a rate limit, we can retry after a short delay.
|
||||
wait_delay += 10;
|
||||
if wait_delay > 60 {
|
||||
wait_delay = 60;
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_secs(wait_delay)).await;
|
||||
continue; // Retry the upload.
|
||||
}
|
||||
Err(err) => {
|
||||
error!("Failed to upload to backend: {:?}", err);
|
||||
return Err(err.into())
|
||||
},
|
||||
}
|
||||
} // End loop
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
}
|
||||
|
||||
let updated = update_task_status(UpdateTaskArgs {
|
||||
db: task_database.get_connection(),
|
||||
task_id: &local_task.id,
|
||||
status: TaskStatus::CompleteSuccess,
|
||||
}).await?;
|
||||
|
||||
if !updated {
|
||||
continue; // If anything breaks with queries, don't spam events.
|
||||
}
|
||||
|
||||
let event = GenerationCompleteEvent {
|
||||
//media_file_token: result.media_file_token,
|
||||
action: Some(GenerationAction::GenerateImage),
|
||||
service: GenerationServiceProvider::Midjourney,
|
||||
model: None,
|
||||
};
|
||||
|
||||
if let Err(err) = event.send(&app_handle) {
|
||||
error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod midjourney_long_polling_thread;
|
||||
@@ -0,0 +1,35 @@
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
use errors::AnyhowResult;
|
||||
use midjourney_client::utils::get_image_url::get_image_url;
|
||||
use std::fs::File;
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use url::Url;
|
||||
use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
|
||||
|
||||
pub async fn download_midjourney_image(
|
||||
image_downloader: &ImageDownloaderClient,
|
||||
job_id: &str,
|
||||
index: u8,
|
||||
app_data_root: &AppDataRoot
|
||||
) -> AnyhowResult<PathBuf> {
|
||||
let url = get_image_url(job_id, index)?;
|
||||
let parsed_url = Url::parse(&url)?;
|
||||
|
||||
//let response = reqwest::get(&url).await?;
|
||||
//let image_bytes = response.bytes().await?;
|
||||
|
||||
let image_bytes = image_downloader.download_image(job_id, index).await?;
|
||||
|
||||
let ext = parsed_url.path().split(".").last().unwrap_or("png");
|
||||
|
||||
let tempdir = app_data_root.temp_dir().path();
|
||||
let download_filename = format!("{}_{}.{}", job_id, index, ext);
|
||||
let download_path = tempdir.join(download_filename);
|
||||
|
||||
let mut file = File::create(&download_path)?;
|
||||
file.write_all(&image_bytes)?;
|
||||
|
||||
Ok(download_path)
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod download_midjourney_image;
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
use cookie_store::cookie_store::CookieStore;
|
||||
use errors::AnyhowResult;
|
||||
use once_cell::sync::Lazy;
|
||||
use reqwest::Url;
|
||||
use std::collections::HashSet;
|
||||
use tauri::webview::Cookie;
|
||||
use tauri::WebviewWindow;
|
||||
|
||||
static WWW_COOKIE_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse("https://www.midjourney.com").expect("URL should parse")
|
||||
});
|
||||
|
||||
static ROOT_COOKIE_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse("https://midjourney.com").expect("URL should parse")
|
||||
});
|
||||
|
||||
pub fn extract_midjourney_webview_cookies(webview: &WebviewWindow) -> AnyhowResult<CookieStore> {
|
||||
let mut cookie_store = CookieStore::empty();
|
||||
let cookies = get_all_midjourney_cookies(webview)?;
|
||||
for cookie in cookies.iter() {
|
||||
cookie_store.add_cookie_name_and_value(
|
||||
cookie.name().to_string(),
|
||||
cookie.value().to_string(),
|
||||
);
|
||||
}
|
||||
Ok(cookie_store)
|
||||
}
|
||||
|
||||
fn get_all_midjourney_cookies(webview: &WebviewWindow) -> AnyhowResult<Vec<Cookie>> {
|
||||
// NB: WWW domain contains the auth/session cookies.
|
||||
// Root domain contains some other cookies, such as Cloudflare.
|
||||
let www_cookies = webview.cookies_for_url(WWW_COOKIE_URL.clone())?;
|
||||
let root_cookies = webview.cookies_for_url(ROOT_COOKIE_URL.clone())?;
|
||||
|
||||
let mut all_cookies = www_cookies;
|
||||
let mut cookie_names = HashSet::new();
|
||||
|
||||
for cookie in root_cookies.iter() {
|
||||
if !cookie_names.contains(cookie.name()) {
|
||||
cookie_names.insert(cookie.name().to_string());
|
||||
all_cookies.push(cookie.clone());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_cookies)
|
||||
}
|
||||
+185
@@ -0,0 +1,185 @@
|
||||
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
|
||||
use crate::core::events::functional_events::refresh_account_state_event::RefreshAccountStateEvent;
|
||||
use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider};
|
||||
use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
|
||||
use crate::core::events::sendable_event_trait::SendableEvent;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::utils::window::get_webview_window_hostname::get_webview_window_hostname;
|
||||
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
|
||||
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
|
||||
use crate::services::midjourney::windows::extract_midjourney_webview_cookies::extract_midjourney_webview_cookies;
|
||||
use crate::services::midjourney::windows::open_midjourney_login_window::MIDJOURNEY_LOGIN_WINDOW_NAME;
|
||||
use crate::services::sora::events::sora_login_success_event::SoraLoginSuccessEvent;
|
||||
use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::services::sora::windows::sora_login_window::extract_sora_webview_cookies::extract_sora_webview_cookies;
|
||||
use anyhow::anyhow;
|
||||
use cookie_store::cookie_store::CookieStore;
|
||||
use enums::common::generation_provider::GenerationProvider;
|
||||
use errors::AnyhowResult;
|
||||
use log::{error, info};
|
||||
use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
|
||||
use midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
|
||||
use midjourney_client::recipes::get_user_info::{get_user_info, GetUserInfoRequest};
|
||||
use openai_sora_client::creds::sora_credential_set::SoraCredentialSet;
|
||||
use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
|
||||
use openai_sora_client::utils::has_session_cookie::{has_session_cookie, SessionCookiePresence};
|
||||
use tauri::{AppHandle, Manager, WebviewWindow};
|
||||
|
||||
pub async fn midjourney_login_window_thread(
|
||||
app: AppHandle,
|
||||
app_data_root: AppDataRoot,
|
||||
mj_creds_manager: MidjourneyCredentialManager,
|
||||
) {
|
||||
let mut visited_login = false;
|
||||
|
||||
loop {
|
||||
let login_webview_window = match app.get_webview_window(MIDJOURNEY_LOGIN_WINDOW_NAME) {
|
||||
Some(webview) => webview,
|
||||
None => {
|
||||
info!("Exit midjourney login thread.");
|
||||
return; // NB: Only exit if we don't have the webview.
|
||||
}
|
||||
};
|
||||
|
||||
let result = check_login_window(
|
||||
&app,
|
||||
&login_webview_window,
|
||||
&app_data_root,
|
||||
&mj_creds_manager,
|
||||
&mut visited_login,
|
||||
).await;
|
||||
|
||||
match result {
|
||||
Err(err) => {
|
||||
error!("Error checking login window: {:?}", err);
|
||||
}
|
||||
Ok(false) => {} // Continue iteration and try again...
|
||||
Ok(true) => {
|
||||
info!("Successfully saved cookies from login window. Closing.");
|
||||
if let Err(err) = login_webview_window.close() {
|
||||
error!("Error closing login window: {:?}", err);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if we can exit.
|
||||
async fn check_login_window(
|
||||
app_handle: &AppHandle,
|
||||
webview_window: &WebviewWindow,
|
||||
app_data_root: &AppDataRoot,
|
||||
mj_creds_manager: &MidjourneyCredentialManager,
|
||||
visited_login: &mut bool,
|
||||
) -> AnyhowResult<bool> {
|
||||
|
||||
/* Login flow looks like this:
|
||||
|
||||
1. Start: https://www.midjourney.com/auth/signin
|
||||
2. SSO / Google Login (etc): https://accounts.google.com/v3/signin/challenge/pwd?[query]
|
||||
3. Done / Landing: https://www.midjourney.com/...
|
||||
*/
|
||||
|
||||
let hostname = get_webview_window_hostname(webview_window)?;
|
||||
|
||||
let mut maybe_at_destination = false;
|
||||
|
||||
match hostname.as_str() {
|
||||
"www.midjourney.com" |
|
||||
"midjourney.com"
|
||||
=> {
|
||||
maybe_at_destination = true;
|
||||
}
|
||||
// chatgpt.com/auth is also an auth domain
|
||||
"auth.openai.com" |
|
||||
"accounts.google.com" |
|
||||
"accounts.youtube.com" |
|
||||
"login.live.com" |
|
||||
"appleid.apple.com"
|
||||
=> {
|
||||
// NB: We're in auth flow.
|
||||
info!("Midjourney webview is in auth flow; hostname `{}`.", hostname);
|
||||
*visited_login = true;
|
||||
return Ok(false)
|
||||
}
|
||||
_ => {}, // We just don't know...
|
||||
}
|
||||
|
||||
let cookie_store = extract_midjourney_webview_cookies(webview_window)?;
|
||||
|
||||
let maybe_has_auth_cookies = cookie_store_has_auth_cookies(&cookie_store);
|
||||
let maybe_has_enough_cookies = cookie_store.len() > 6;
|
||||
let maybe_completed_login_cycle = *visited_login && maybe_at_destination;
|
||||
|
||||
// Misc cookies without login cookies are ~1055 length
|
||||
// AUTH_I is ~1500 length
|
||||
// AUTH_R is ~500 length
|
||||
let maybe_has_big_enough_cookie = cookie_store.calculate_approx_cookie_character_length() > 2100;
|
||||
|
||||
let mut heuristic_count = 0;
|
||||
|
||||
if maybe_has_auth_cookies {
|
||||
heuristic_count += 1;
|
||||
}
|
||||
if maybe_has_enough_cookies {
|
||||
heuristic_count += 1;
|
||||
}
|
||||
if maybe_completed_login_cycle {
|
||||
heuristic_count += 1;
|
||||
}
|
||||
if maybe_has_big_enough_cookie {
|
||||
heuristic_count += 1;
|
||||
}
|
||||
|
||||
if heuristic_count < 3 {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
info!("Heuristic count is ({}); we're likely done.", heuristic_count);
|
||||
|
||||
info!("Current cookies (len {}): {:?}", cookie_store.len(), cookie_store.to_cookie_string());
|
||||
|
||||
let response = get_user_info(GetUserInfoRequest {
|
||||
hostname: MidjourneyHostname::Standard,
|
||||
cookie_header: cookie_store.to_cookie_string(),
|
||||
}).await;
|
||||
|
||||
match response {
|
||||
Err(err) => {
|
||||
error!("Error getting midjourney user info: {:?}", err);
|
||||
// NB: Fall through. Allow us to update the cookies.
|
||||
}
|
||||
Ok(user_info) => {
|
||||
info!("Got midjourney user info: {:?}", user_info);
|
||||
let user_info = MidjourneyUserInfo::from_api_response(user_info);
|
||||
|
||||
if let Err(err) = mj_creds_manager.replace_user_info(user_info) {
|
||||
error!("Error saving midjourney user info: {:?}", err);
|
||||
// NB: Fall through. Allow us to update the cookies.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mj_creds_manager.replace_cookie_store(cookie_store)?;
|
||||
|
||||
let result = mj_creds_manager.persist_to_disk();
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("Error persisting midjourney cookies to disk: {:?}", err);
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let event = RefreshAccountStateEvent {
|
||||
provider: Some(GenerationProvider::Midjourney),
|
||||
};
|
||||
|
||||
if let Err(err) = event.send(&app_handle) {
|
||||
error!("Failed to send RefreshAccountStateEvent: {:?}", err); // Fail open
|
||||
}
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
pub mod extract_midjourney_webview_cookies;
|
||||
pub mod midjourney_login_window_thread;
|
||||
pub mod open_midjourney_login_window;
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
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;
|
||||
use crate::services::midjourney::windows::midjourney_login_window_thread::midjourney_login_window_thread;
|
||||
use anyhow::anyhow;
|
||||
use errors::AnyhowResult;
|
||||
use once_cell::sync::Lazy;
|
||||
use reqwest::Url;
|
||||
use std::time::Duration;
|
||||
use tauri::{AppHandle, Manager, WebviewUrl, WebviewWindowBuilder};
|
||||
|
||||
/// Name of the window
|
||||
pub (super) const MIDJOURNEY_LOGIN_WINDOW_NAME: &str = "midjourney_login_window";
|
||||
|
||||
pub (super) const MIDJOURNEY_START_URL_STR: &str = "https://google.com";
|
||||
|
||||
pub (super) static START_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse(MIDJOURNEY_START_URL_STR).expect("URL should parse")
|
||||
});
|
||||
|
||||
pub (super) const MIDJOURNEY_HOMEPAGE_URL_STR: &str = "https://www.midjourney.com/";
|
||||
|
||||
pub (super) static MIDJOURNEY_HOMEPAGE_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse(MIDJOURNEY_HOMEPAGE_URL_STR).expect("URL should parse")
|
||||
});
|
||||
|
||||
pub (super) const MIDJOURNEY_LOGIN_URL_STR: &str = "https://www.midjourney.com/auth/signin";
|
||||
|
||||
pub (super) static MIDJOURNEY_LOGIN_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse(MIDJOURNEY_LOGIN_URL_STR).expect("URL should parse")
|
||||
});
|
||||
|
||||
pub async fn open_midjourney_login_window(
|
||||
app: &AppHandle,
|
||||
app_data_root: &AppDataRoot,
|
||||
mj_creds_manager: &MidjourneyCredentialManager,
|
||||
) -> AnyhowResult<()> {
|
||||
if app.get_window(MIDJOURNEY_LOGIN_WINDOW_NAME).is_some() {
|
||||
return Err(anyhow!("Login window already open"));
|
||||
}
|
||||
|
||||
let url = WebviewUrl::External(START_URL.clone());
|
||||
|
||||
let window = WebviewWindowBuilder::new(app, MIDJOURNEY_LOGIN_WINDOW_NAME, url)
|
||||
//.user_agent(openai_sora_client::credentials::USER_AGENT)
|
||||
.always_on_top(false)
|
||||
.title("Login to Midjourney")
|
||||
.center()
|
||||
.resizable(true)
|
||||
.visible(true)
|
||||
.closable(true)
|
||||
.min_inner_size(200.0, 800.0)
|
||||
.focused(true)
|
||||
.devtools(true)
|
||||
.build()?;
|
||||
|
||||
let webview = window.get_webview(MIDJOURNEY_LOGIN_WINDOW_NAME)
|
||||
.ok_or_else(|| anyhow!("no webview found"))?;
|
||||
|
||||
clear_all_webview_cookies(&webview)?;
|
||||
|
||||
webview.navigate(MIDJOURNEY_HOMEPAGE_URL.clone())?;
|
||||
|
||||
// NB: We're starting to get Cloudflare protection screens. Let's try to avoid.
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
|
||||
webview.navigate(MIDJOURNEY_LOGIN_URL.clone())?;
|
||||
|
||||
let app_handle = app.clone();
|
||||
let app_data_root = app_data_root.clone();
|
||||
let mj_creds_manager= mj_creds_manager.clone();
|
||||
|
||||
let _ = tauri::async_runtime::spawn(async move {
|
||||
midjourney_login_window_thread(app_handle, app_data_root, mj_creds_manager).await;
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod fal;
|
||||
pub mod midjourney;
|
||||
pub mod sora;
|
||||
pub mod storyteller;
|
||||
@@ -140,6 +140,7 @@ async fn polling_loop(
|
||||
maybe_creds: Some(&creds),
|
||||
path: download_path,
|
||||
is_intermediate_system_file: false,
|
||||
maybe_prompt_token: None, // TODO: this should be added soon.
|
||||
}).await?;
|
||||
|
||||
info!("Uploaded to API backend: {:?}", result.media_file_token);
|
||||
|
||||
+7
-7
@@ -10,23 +10,23 @@ use std::time::Duration;
|
||||
use tauri::{AppHandle, Manager, WebviewUrl, WebviewWindowBuilder};
|
||||
|
||||
/// Name of the window
|
||||
pub const LOGIN_WINDOW_NAME: &str = "login_window";
|
||||
pub (super) const LOGIN_WINDOW_NAME: &str = "login_window";
|
||||
|
||||
pub const START_URL_STR: &str = "https://google.com";
|
||||
pub (super) const START_URL_STR: &str = "https://google.com";
|
||||
|
||||
pub static START_URL: Lazy<Url> = Lazy::new(|| {
|
||||
pub (super) static START_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse(START_URL_STR).expect("URL should parse")
|
||||
});
|
||||
|
||||
pub const SORA_HOMEPAGE_URL_STR: &str = "https://sora.com/";
|
||||
pub (super) const SORA_HOMEPAGE_URL_STR: &str = "https://sora.com/";
|
||||
|
||||
pub static SORA_HOMEPAGE_URL : Lazy<Url> = Lazy::new(|| {
|
||||
pub (super) static SORA_HOMEPAGE_URL : Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse(SORA_HOMEPAGE_URL_STR).expect("URL should parse")
|
||||
});
|
||||
|
||||
pub const SORA_LOGIN_URL_STR: &str = "https://chatgpt.com/auth/login?next=%2Fsora%2F";
|
||||
pub (super) const SORA_LOGIN_URL_STR: &str = "https://chatgpt.com/auth/login?next=%2Fsora%2F";
|
||||
|
||||
pub static SORA_LOGIN_URL: Lazy<Url> = Lazy::new(|| {
|
||||
pub (super) static SORA_LOGIN_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse(SORA_LOGIN_URL_STR).expect("URL should parse")
|
||||
});
|
||||
|
||||
|
||||
+2
-9
@@ -1,5 +1,6 @@
|
||||
use crate::core::events::sendable_event_trait::SendableEvent;
|
||||
use crate::core::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::core::utils::window::get_webview_window_hostname::get_webview_window_hostname;
|
||||
use crate::services::sora::events::sora_login_success_event::SoraLoginSuccessEvent;
|
||||
use crate::services::sora::state::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::services::sora::windows::sora_login_window::extract_sora_webview_cookies::extract_sora_webview_cookies;
|
||||
@@ -73,7 +74,7 @@ async fn check_login_window(
|
||||
- Check if we're on the correct domain, if not exit? (or inverse, that we're not in login flow)
|
||||
*/
|
||||
|
||||
let hostname = get_hostname(webview_window)?;
|
||||
let hostname = get_webview_window_hostname(webview_window)?;
|
||||
|
||||
let mut maybe_at_destination = false;
|
||||
|
||||
@@ -138,11 +139,3 @@ async fn check_login_window(
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn get_hostname(webview: &WebviewWindow) -> AnyhowResult<String> {
|
||||
let url = webview.url()?;
|
||||
let url_hostname= url.host()
|
||||
.ok_or(anyhow!("no host in url"))?
|
||||
.to_string();
|
||||
Ok(url_hostname)
|
||||
}
|
||||
|
||||
+1
-1
@@ -49,7 +49,7 @@ impl StorytellerCredentialHolder {
|
||||
match self.credentials.read() {
|
||||
Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)),
|
||||
Ok(creds) => match &*creds {
|
||||
None => Err(anyhow!("Credentials not set")),
|
||||
None => Err(anyhow!("Storyteller credentials not set")),
|
||||
Some(creds) => Ok(creds.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,6 @@
|
||||
"frontendDist": "../../../frontend/apps/editor3d/dist",
|
||||
"beforeDevCommand": "echo 'no dev'",
|
||||
"devUrl": "http://localhost:5173",
|
||||
"beforeBuildCommand": "nx run editor3d:build"
|
||||
"beforeBuildCommand": "npx nx run editor3d:build"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,9 +3,16 @@ use std::fmt::{Display, Formatter};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum CloudflareError {
|
||||
/// Cloudflare returned a 301 Moved Permanently response.
|
||||
/// TODO: This needs to include "location" and "set-cookie" headers.
|
||||
MovedPermanently301,
|
||||
|
||||
/// Cloudflare wants to verify the request with a CAPTCHA challenge.
|
||||
ChallengeInterstitial403,
|
||||
|
||||
|
||||
/// Cloudflare returned a 502 Bad Gateway response.
|
||||
BadGateway502,
|
||||
|
||||
/// Cloudflare could not form a connection to the backend server.
|
||||
GatewayTimeout504,
|
||||
|
||||
@@ -18,9 +25,15 @@ impl Error for CloudflareError {}
|
||||
impl Display for CloudflareError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::MovedPermanently301 => {
|
||||
write!(f, "Cloudflare Moved Permanently (301)")
|
||||
}
|
||||
Self::ChallengeInterstitial403 => {
|
||||
write!(f, "Cloudflare Challenge Interstitial (403); Cloudflare wants to verify the request with a CAPTCHA challenge.")
|
||||
}
|
||||
Self::BadGateway502 => {
|
||||
write!(f, "Cloudflare Bad Gateway (502); This is likely a backend server issue.")
|
||||
}
|
||||
Self::GatewayTimeout504 => {
|
||||
write!(f, "Cloudflare Gateway Timeout (504); This is likely a backend server issue.")
|
||||
}
|
||||
|
||||
@@ -22,11 +22,28 @@ pub fn filter_cloudflare_errors(status_code: u16, body: &str) -> Result<(), Clou
|
||||
let is_cloudflare = body.contains("cloudflare")
|
||||
|| body.contains("Cloudflare");
|
||||
|
||||
if status_code == 301 {
|
||||
/*
|
||||
<html>
|
||||
<head><title>301 Moved Permanently</title></head>
|
||||
<body>
|
||||
<center><h1>301 Moved Permanently</h1></center>
|
||||
<hr><center>cloudflare</center>
|
||||
</body>
|
||||
</html>
|
||||
*/
|
||||
if is_cloudflare {
|
||||
// TODO: Include location header.
|
||||
return Err(CloudflareError::MovedPermanently301);
|
||||
}
|
||||
}
|
||||
|
||||
// let is_cloudflare = body.contains("cloudflare.com")
|
||||
// || body.contains("Cloudflare Ray ID");
|
||||
|
||||
if is_cloudflare {
|
||||
match status_code {
|
||||
502 => return Err(CloudflareError::BadGateway502),
|
||||
504 => return Err(CloudflareError::GatewayTimeout504),
|
||||
524 => return Err(CloudflareError::TimeoutOccurred524),
|
||||
_ => {}, // Fall-through
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
[package]
|
||||
name = "cookie_store"
|
||||
edition = "2024"
|
||||
version = "0.0.1"
|
||||
publish = false
|
||||
|
||||
[package.metadata.cargo-udeps.ignore]
|
||||
normal = ["workspace-build-acceleration"]
|
||||
|
||||
[lib]
|
||||
name = "cookie_store"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[dependencies]
|
||||
# Workspace hack (faster compile times)
|
||||
workspace-build-acceleration.workspace = true
|
||||
|
||||
# Internal
|
||||
errors.workspace = true
|
||||
|
||||
# External
|
||||
serde.workspace = true
|
||||
serde_derive.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
# None yet
|
||||
@@ -0,0 +1,34 @@
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Cookie {
|
||||
pub name: String,
|
||||
pub value: String,
|
||||
|
||||
// TODO: Later.
|
||||
// pub domain: String,
|
||||
// pub path: String,
|
||||
// pub secure: bool,
|
||||
// pub http_only: bool,
|
||||
// pub expires: Option<String>,
|
||||
}
|
||||
|
||||
impl Cookie {
|
||||
pub fn new(name: String, value: String) -> Self {
|
||||
Self { name, value }
|
||||
}
|
||||
|
||||
pub fn new_from_str(name: &str, value: &str) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
value: value.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
pub fn value(&self) -> &str {
|
||||
&self.value
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
use crate::cookie::Cookie;
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub struct CookieSet {
|
||||
// TODO: This doesn't yet handle partial paths, subddomain rules, etc.
|
||||
pub cookies: HashMap<String, Cookie>,
|
||||
}
|
||||
|
||||
impl CookieSet {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
cookies: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_cookie(&mut self, cookie: Cookie) {
|
||||
self.cookies.insert(cookie.name.clone(), cookie);
|
||||
}
|
||||
|
||||
pub fn get_cookie(&self, name: &str) -> Option<&Cookie> {
|
||||
self.cookies.get(name)
|
||||
}
|
||||
|
||||
pub fn remove_cookie(&mut self, name: &str) {
|
||||
self.cookies.remove(name);
|
||||
}
|
||||
|
||||
pub fn to_cookie_string(&self) -> String {
|
||||
self
|
||||
.cookies
|
||||
.values()
|
||||
.map(|cookie| format!("{}={}", cookie.name, cookie.value))
|
||||
.collect::<Vec<String>>()
|
||||
.join("; ")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
use crate::cookie::Cookie;
|
||||
use crate::serialized_cookie_store::{SerializableCookieStore, SerializedCookie};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CookieStore {
|
||||
// TODO: This doesn't yet handle partial paths, subdomain rules, etc.
|
||||
pub cookies: HashMap<String, Cookie>,
|
||||
}
|
||||
|
||||
impl CookieStore {
|
||||
pub fn empty() -> Self {
|
||||
Self {
|
||||
cookies: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_cookie(&mut self, cookie: Cookie) {
|
||||
self.cookies.insert(cookie.name.clone(), cookie);
|
||||
}
|
||||
|
||||
pub fn add_cookie_name_and_value(&mut self, name: String, value: String) {
|
||||
self.add_cookie(Cookie {
|
||||
name,
|
||||
value,
|
||||
});
|
||||
}
|
||||
|
||||
pub fn get_cookie(&self, name: &str) -> Option<&Cookie> {
|
||||
self.cookies.get(name)
|
||||
}
|
||||
|
||||
pub fn has_cookie(&self, name: &str) -> bool {
|
||||
self.cookies.contains_key(name)
|
||||
}
|
||||
|
||||
pub fn remove_cookie(&mut self, name: &str) {
|
||||
self.cookies.remove(name);
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.cookies.len()
|
||||
}
|
||||
|
||||
pub fn clear_all(&mut self) {
|
||||
self.cookies.clear();
|
||||
}
|
||||
|
||||
/// NB: Just use this as a heuristic, and do not call it in a loop.
|
||||
pub fn calculate_approx_cookie_character_length(&self) -> usize {
|
||||
const COOKIE_SEP_LENGTH: usize = 2; // '=' and ';'
|
||||
self.cookies
|
||||
.values()
|
||||
.map(|cookie| cookie.name.len() + cookie.value.len() + COOKIE_SEP_LENGTH)
|
||||
.sum()
|
||||
}
|
||||
|
||||
pub fn to_cookie_string(&self) -> String {
|
||||
self
|
||||
.cookies
|
||||
.values()
|
||||
.map(|cookie| format!("{}={}", cookie.name, cookie.value))
|
||||
.collect::<Vec<String>>()
|
||||
.join("; ")
|
||||
}
|
||||
|
||||
pub fn to_serializable(&self) -> SerializableCookieStore {
|
||||
let cookies = self
|
||||
.cookies
|
||||
.values()
|
||||
.map(|cookie| SerializedCookie {
|
||||
name: cookie.name.clone(),
|
||||
value: cookie.value.clone(),
|
||||
})
|
||||
.collect();
|
||||
SerializableCookieStore { cookies }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
//! cookie_store
|
||||
//!
|
||||
//! Lightweight cookie jar
|
||||
//!
|
||||
|
||||
// Never allow these
|
||||
#![forbid(private_bounds)]
|
||||
#![forbid(private_interfaces)]
|
||||
#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things
|
||||
|
||||
// Okay to toggle
|
||||
#![forbid(unreachable_patterns)]
|
||||
#![forbid(unused_imports)]
|
||||
#![forbid(unused_mut)]
|
||||
#![forbid(unused_variables)]
|
||||
|
||||
// Always allow
|
||||
#![allow(dead_code)]
|
||||
#![allow(non_snake_case)]
|
||||
|
||||
pub mod cookie;
|
||||
pub mod cookie_set;
|
||||
pub mod cookie_store;
|
||||
pub mod serialized_cookie_store;
|
||||
@@ -0,0 +1,66 @@
|
||||
use crate::cookie::Cookie;
|
||||
use crate::cookie_store::CookieStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::error::Error;
|
||||
use std::fmt::{Debug, Display, Formatter};
|
||||
use std::path::Path;
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct SerializableCookieStore {
|
||||
pub (crate) cookies: Vec<SerializedCookie>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub (crate) struct SerializedCookie {
|
||||
pub (crate) name: String,
|
||||
pub (crate) value: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SerializableCookieStoreError {
|
||||
IoError(std::io::Error),
|
||||
SerializationError(serde_json::Error),
|
||||
}
|
||||
|
||||
impl Error for SerializableCookieStoreError {}
|
||||
|
||||
impl Display for SerializableCookieStoreError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::IoError(err) => write!(f, "IO error: {}", err),
|
||||
Self::SerializationError(err) => write!(f, "Serialization error: {}", err),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SerializableCookieStore {
|
||||
pub fn read_from_file<P: AsRef<Path>>(
|
||||
file_path: P,
|
||||
) -> Result<Self, SerializableCookieStoreError> {
|
||||
let content = std::fs::read_to_string(file_path)
|
||||
.map_err(|err| SerializableCookieStoreError::IoError(err))?;
|
||||
serde_json::from_str(&content)
|
||||
.map_err(|err| SerializableCookieStoreError::SerializationError(err))
|
||||
}
|
||||
|
||||
pub fn write_to_file<P: AsRef<Path>>(
|
||||
&self,
|
||||
file_path: P,
|
||||
) -> Result<(), SerializableCookieStoreError> {
|
||||
let serialized = serde_json::to_string(self)
|
||||
.map_err(|err| SerializableCookieStoreError::SerializationError(err))?;
|
||||
std::fs::write(file_path, serialized)
|
||||
.map_err(|err| SerializableCookieStoreError::IoError(err))
|
||||
}
|
||||
|
||||
pub fn to_cookie_store(&self) -> CookieStore {
|
||||
let mut store = CookieStore::empty();
|
||||
for serialized_cookie in &self.cookies {
|
||||
store.add_cookie(Cookie::new_from_str(
|
||||
&serialized_cookie.name,
|
||||
&serialized_cookie.value
|
||||
));
|
||||
}
|
||||
store
|
||||
}
|
||||
}
|
||||
@@ -31,5 +31,6 @@ pub mod file_size;
|
||||
pub mod filename_concat;
|
||||
pub mod is_filesystem_full_error;
|
||||
pub mod path_to_string;
|
||||
pub mod read_to_trimmed_string;
|
||||
pub mod recursively_find_file_by_name;
|
||||
pub mod rename_across_devices;
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user