diff --git a/Cargo.lock b/Cargo.lock index c6ce971c..9c8483c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -214,6 +214,7 @@ version = "0.0.1" dependencies = [ "anyhow", "artcraft_client", + "async-trait", "base64 0.22.1", "browser_emulation", "build_metadata", @@ -423,13 +424,13 @@ checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" [[package]] name = "async-trait" -version = "0.1.80" +version = "0.1.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6fa2087f2753a7da8cc1c0dbfcf89579dd57458e36769de5ac750b4671737ca" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.5", ] [[package]] @@ -5658,8 +5659,10 @@ version = "0.0.1" dependencies = [ "anyhow", "artcraft_client", + "async-trait", "browser_emulation", "fal_client", + "file_hashing", "gmicloud_client", "grok_api_client", "grok_consumer_client", @@ -6887,6 +6890,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" diff --git a/Cargo.toml b/Cargo.toml index 6fcdb876..6266636d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -139,6 +139,7 @@ worldlabs_consumer_client = { path = "crates/api_clients/worldlabs_consumer_clie anyhow = "1.0.102" assertor = "0.0.3" base64 = "0.22.1" +async-trait = "0.1.89" blake3 = "1.5" bstr = "1.12.1" byteorder = "1.5.0" diff --git a/crates/api_clients/higgsfield_client/src/session/upload_media.rs b/crates/api_clients/higgsfield_client/src/session/upload_media.rs index b3e3d7a8..46f5f88c 100644 --- a/crates/api_clients/higgsfield_client/src/session/upload_media.rs +++ b/crates/api_clients/higgsfield_client/src/session/upload_media.rs @@ -232,6 +232,16 @@ impl HiggsfieldSession { } impl HiggsfieldSession { + /// One status read of an uploaded image or clip (`GET /fnf/media/{id}` / + /// `/fnf/video/{id}`). Used to confirm a previously uploaded reference + /// still exists (and whether its IP check ran) before reusing its id. + pub async fn get_media_status(&self, family: MediaStatusFamily, media_id: &MediaId) -> Result { + self.with_auth(|auth| { + let request = GetMediaStatusRequest::new(family, media_id.clone()); + async move { get_media_status(GetMediaStatusArgs { request, auth: &auth, host: self.api_host() }).await } + }).await + } + /// Poll the upload's status (`GET /fnf/media/{id}` for images, /// `/fnf/video/{id}` for clips) until its IP check has finished. Only /// completes if the check was requested (`force_ip_check` on confirm); diff --git a/crates/artcraftx/Cargo.toml b/crates/artcraftx/Cargo.toml index 264dfc87..9fb64c08 100644 --- a/crates/artcraftx/Cargo.toml +++ b/crates/artcraftx/Cargo.toml @@ -56,6 +56,7 @@ worldlabs_consumer_client.workspace = true # External anyhow.workspace = true +async-trait.workspace = true base64.workspace = true bytes = "1.10.1" chrono.workspace = true diff --git a/crates/artcraftx/src/commands/generate/common/higgsfield_generation.rs b/crates/artcraftx/src/commands/generate/common/higgsfield_generation.rs index 3ac62e13..17189de6 100644 --- a/crates/artcraftx/src/commands/generate/common/higgsfield_generation.rs +++ b/crates/artcraftx/src/commands/generate/common/higgsfield_generation.rs @@ -8,6 +8,7 @@ use std::collections::HashMap; use artcraft_client::utils::api_host::ApiHost; use log::{info, warn}; +use router::api::asset_upload_cache::AssetUploadCache; use router::client::router_client::RouterClient; use router::client::router_higgsfield_client::RouterHiggsfieldClient; use router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; @@ -53,6 +54,7 @@ pub async fn send_higgsfield_image_request( builder: GenerateImageRequestBuilder, client: &RouterClient, media_url_map: &HashMap, + upload_cache: &dyn AssetUploadCache, ) -> Result { let request = match builder.build2().map_err(|err| { warn!("Could not build Higgsfield image request: {:?}", err); @@ -64,6 +66,7 @@ pub async fn send_higgsfield_image_request( let context = ImageGenerationDraftContext { client: Some(client), media_file_to_artcraft_url_map: Some(media_url_map), + asset_upload_cache: Some(upload_cache), }; draft.finalize(context).await.map_err(|err| { warn!("Could not upload references to Higgsfield: {:?}", err); @@ -84,6 +87,7 @@ pub async fn send_higgsfield_video_request( builder: GenerateVideoRequestBuilder, client: &RouterClient, media_url_map: &HashMap, + upload_cache: &dyn AssetUploadCache, ) -> Result { let request = match builder.build2().map_err(|err| { warn!("Could not build Higgsfield video request: {:?}", err); @@ -96,6 +100,7 @@ pub async fn send_higgsfield_video_request( client: Some(client), media_file_to_artcraft_url_map: Some(media_url_map), character_token_to_kinovi_id_map: None, + asset_upload_cache: Some(upload_cache), }; draft.finalize(context).await.map_err(|err| { warn!("Could not upload media to Higgsfield: {:?}", err); diff --git a/crates/artcraftx/src/commands/generate/common/media_source_conversion.rs b/crates/artcraftx/src/commands/generate/common/media_source_conversion.rs index 819b6fdc..cc0a3919 100644 --- a/crates/artcraftx/src/commands/generate/common/media_source_conversion.rs +++ b/crates/artcraftx/src/commands/generate/common/media_source_conversion.rs @@ -31,6 +31,7 @@ use uuid_utils::uuid::generate_random_uuid; use crate::commands::generate::common::tauri_media_source::TauriMediaSource; use crate::commands::generate::generate_error::GenerateError; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; /// Which ArtCraft upload endpoint a source goes through. #[derive(Copy, Clone, Debug, PartialEq, Eq)] @@ -78,22 +79,29 @@ pub fn collect_source_tokens<'a>(sources: impl IntoIterator, api_host: &ApiHost, + ledger: &AssetUploadLedger, ) -> Result { + let maybe_cache = ledger.for_artcraft_session(maybe_creds); match source { TauriMediaSource::MediaFileToken { token } => Ok(token), - TauriMediaSource::LocalPath { path } => upload_path_to_artcraft(&path, kind, maybe_creds, api_host).await, + TauriMediaSource::LocalPath { path } => match &maybe_cache { + Some(cache) => cache.artcraft_token_for_file(api_host, &path, || upload_path_to_artcraft(&path, kind, maybe_creds, api_host)).await, + None => upload_path_to_artcraft(&path, kind, maybe_creds, api_host).await, + }, TauriMediaSource::Bytes { bytes, file_name } => { - // The upload endpoints are multipart-from-file; stage the bytes. - let staged = stage_bytes_to_temp_file(&bytes, file_name.as_deref())?; - let result = upload_path_to_artcraft(&staged, kind, maybe_creds, api_host).await; - std::fs::remove_file(&staged).ok(); - result + let upload = || upload_bytes_to_artcraft(&bytes, file_name.as_deref(), kind, maybe_creds, api_host); + match &maybe_cache { + Some(cache) => cache.artcraft_token_for_bytes(api_host, &bytes, upload).await, + None => upload().await, + } } } } @@ -104,6 +112,7 @@ pub async fn sources_to_artcraft_tokens( kind: ArtcraftMediaKind, maybe_creds: Option<&StorytellerCredentialSet>, api_host: &ApiHost, + ledger: &AssetUploadLedger, ) -> Result>, GenerateError> { let sources = match maybe_sources { None => return Ok(None), @@ -112,7 +121,7 @@ pub async fn sources_to_artcraft_tokens( }; let mut tokens = Vec::with_capacity(sources.len()); for source in sources { - tokens.push(source_to_artcraft_token(source, kind, maybe_creds, api_host).await?); + tokens.push(source_to_artcraft_token(source, kind, maybe_creds, api_host, ledger).await?); } Ok(Some(tokens)) } @@ -123,10 +132,11 @@ pub async fn maybe_source_to_artcraft_token( kind: ArtcraftMediaKind, maybe_creds: Option<&StorytellerCredentialSet>, api_host: &ApiHost, + ledger: &AssetUploadLedger, ) -> Result, GenerateError> { match maybe_source { None => Ok(None), - Some(source) => Ok(Some(source_to_artcraft_token(source, kind, maybe_creds, api_host).await?)), + Some(source) => Ok(Some(source_to_artcraft_token(source, kind, maybe_creds, api_host, ledger).await?)), } } @@ -138,6 +148,7 @@ pub async fn image_sources_to_fal_urls( maybe_sources: Option>, maybe_creds: Option<&StorytellerCredentialSet>, api_host: &ApiHost, + ledger: &AssetUploadLedger, ) -> Result>, GenerateError> { let sources = match maybe_sources { None => return Ok(None), @@ -146,7 +157,7 @@ pub async fn image_sources_to_fal_urls( }; let mut urls = Vec::with_capacity(sources.len()); for source in sources { - let token = source_to_artcraft_token(source, ArtcraftMediaKind::Image, maybe_creds, api_host).await?; + let token = source_to_artcraft_token(source, ArtcraftMediaKind::Image, maybe_creds, api_host, ledger).await?; let response = get_media_file(api_host, &token).await?; urls.push(response.media_file.media_links.cdn_url.to_string()); } @@ -155,6 +166,21 @@ pub async fn image_sources_to_fal_urls( // ── Private ── +/// The upload endpoints are multipart-from-file; stage the bytes, upload, +/// clean up. +async fn upload_bytes_to_artcraft( + bytes: &[u8], + maybe_file_name: Option<&str>, + kind: ArtcraftMediaKind, + maybe_creds: Option<&StorytellerCredentialSet>, + api_host: &ApiHost, +) -> Result { + let staged = stage_bytes_to_temp_file(bytes, maybe_file_name)?; + let result = upload_path_to_artcraft(&staged, kind, maybe_creds, api_host).await; + std::fs::remove_file(&staged).ok(); + result +} + async fn upload_path_to_artcraft( path: &Path, kind: ArtcraftMediaKind, diff --git a/crates/artcraftx/src/commands/generate/generate_audio/generate_audio_command.rs b/crates/artcraftx/src/commands/generate/generate_audio/generate_audio_command.rs index 8765487f..6f38ca6f 100644 --- a/crates/artcraftx/src/commands/generate/generate_audio/generate_audio_command.rs +++ b/crates/artcraftx/src/commands/generate/generate_audio/generate_audio_command.rs @@ -13,6 +13,8 @@ use crate::events::generation_events::generation_enqueue_success_event::Generati use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::database::local_files_database::LocalFilesDatabase; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::state::database::task_database::TaskDatabase; use log::{error, info}; use tauri::{AppHandle, State}; @@ -22,17 +24,18 @@ pub async fn generate_audio_command( request: TauriGenerateAudioRequest, app: AppHandle, app_data_root: State<'_, AppDataRoot>, - task_database: State<'_, TaskDatabase>, + local_files_database: State<'_, LocalFilesDatabase>, ) -> Response { info!("generate_audio_command called, request: {:?}", request); + let ledger = AssetUploadLedger::new(local_files_database.inner().clone()); let result = handle_request( request, &app_data_root, - &task_database, + &ledger, ).await; match result { @@ -90,13 +93,13 @@ pub async fn generate_audio_command( async fn handle_request( request: TauriGenerateAudioRequest, app_data_root: &AppDataRoot, - task_database: &TaskDatabase, + ledger: &AssetUploadLedger, ) -> Result { // Generation is credential-driven: the request names a stored credential, // and the router dispatches to that credential's service. - let success_event = handle_credential_router(&request, app_data_root).await?; + let success_event = handle_credential_router(&request, app_data_root, ledger).await?; let result = success_event .insert_into_task_database_with_frontend_payload( diff --git a/crates/artcraftx/src/commands/generate/generate_audio/handle_credential_router.rs b/crates/artcraftx/src/commands/generate/generate_audio/handle_credential_router.rs index bf15ecde..2fd86816 100644 --- a/crates/artcraftx/src/commands/generate/generate_audio/handle_credential_router.rs +++ b/crates/artcraftx/src/commands/generate/generate_audio/handle_credential_router.rs @@ -13,6 +13,7 @@ use router::api::image_list_ref::ImageListRef; use crate::commands::utils::api_adapters::models::audio::tauri_audio_model_to_generation_model::tauri_audio_model_to_generation_model; use crate::commands::utils::api_adapters::models::audio::tauri_audio_model_to_router_model::tauri_audio_model_to_router_model; use crate::commands::generate::generate_error::GenerateError; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::commands::generate::task_enqueue_success::TaskEnqueueSuccess; use crate::commands::generate::common::generation_credential::{credential_not_usable, resolve_generation_credential, storyteller_creds_from_credential}; use crate::commands::generate::common::media_source_conversion::{sources_to_artcraft_tokens, ArtcraftMediaKind}; @@ -28,6 +29,7 @@ use crate::state::data_dir::app_data_root::AppDataRoot; pub async fn handle_credential_router( request: &TauriGenerateAudioRequest, app_data_root: &AppDataRoot, + ledger: &AssetUploadLedger, ) -> Result { // Reject unusable media sources (missing local files, empty bytes) before // any provider work. @@ -49,7 +51,7 @@ pub async fn handle_credential_router( GenerationSource::Artcraft | GenerationSource::ArtcraftLocal | GenerationSource::ArtcraftCookies => { - handle_artcraft_credential(request, &credential).await + handle_artcraft_credential(request, &credential, ledger).await } other => Err(credential_not_usable( &credential, @@ -64,6 +66,7 @@ pub async fn handle_credential_router( async fn handle_artcraft_credential( request: &TauriGenerateAudioRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -76,10 +79,10 @@ async fn handle_artcraft_credential( let creds = storyteller_creds_from_credential(credential)?; let audio_references = sources_to_artcraft_tokens( - request.audio_media_sources(), ArtcraftMediaKind::Audio, Some(&creds), &api_host, + request.audio_media_sources(), ArtcraftMediaKind::Audio, Some(&creds), &api_host, ledger, ).await?.map(AudioListRef::MediaFileTokens); let image_references = sources_to_artcraft_tokens( - request.image_media_sources(), ArtcraftMediaKind::Image, Some(&creds), &api_host, + request.image_media_sources(), ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger, ).await?.map(ImageListRef::MediaFileTokens); let client = RouterClient::Artcraft(RouterArtcraftClient::new( diff --git a/crates/artcraftx/src/commands/generate/generate_image/enqueue_image_generation.rs b/crates/artcraftx/src/commands/generate/generate_image/enqueue_image_generation.rs index 163d3c82..b018c30a 100644 --- a/crates/artcraftx/src/commands/generate/generate_image/enqueue_image_generation.rs +++ b/crates/artcraftx/src/commands/generate/generate_image/enqueue_image_generation.rs @@ -45,6 +45,7 @@ use crate::commands::utils::api_adapters::models::image::tauri_image_model_to_ge use crate::commands::utils::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model; use crate::credentials::auth_credential::AuthCredential; use crate::credentials::find_service_credentials::find_first_credential_for_service; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::services::grok::state::grok_websockets::GrokWebsockets; use crate::services::grok::util::refresh_grok_statsig::{refresh_grok_statsig_blocking, GrokStatsigRefreshOutcome}; use crate::services::midjourney::state::midjourney_live_session::MidjourneyLiveSession; @@ -64,6 +65,7 @@ pub async fn enqueue_image_generation( app_data_root: &AppDataRoot, mj_session: &MidjourneyLiveSession, grok_websockets: &GrokWebsockets, + ledger: &AssetUploadLedger, ) -> Result { // Reject unusable media sources (missing local files, empty bytes) before // any provider work. @@ -81,13 +83,13 @@ pub async fn enqueue_image_generation( match credential.service { GenerationSource::Artcraft | GenerationSource::ArtcraftLocal - | GenerationSource::ArtcraftCookies => enqueue_via_artcraft(request, &credential).await, + | GenerationSource::ArtcraftCookies => enqueue_via_artcraft(request, &credential, ledger).await, GenerationSource::FalApi => { let api_key = credential.api_key().ok_or_else(|| { credential_not_usable(&credential, "the FAL credential has no API key") })?; - enqueue_via_fal(request, &api_key.api_key, app_data_root).await + enqueue_via_fal(request, &api_key.api_key, app_data_root, ledger).await } GenerationSource::MidjourneyCookies | GenerationSource::Midjourney => { @@ -99,7 +101,7 @@ pub async fn enqueue_image_generation( } GenerationSource::HiggsfieldCookies | GenerationSource::Higgsfield => { - enqueue_via_higgsfield(request, &credential).await + enqueue_via_higgsfield(request, &credential, ledger).await } other => Err(credential_not_usable( @@ -117,6 +119,7 @@ pub async fn enqueue_image_generation( async fn enqueue_via_artcraft( request: &TauriGenerateImageRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -137,7 +140,7 @@ async fn enqueue_via_artcraft( // the media must reach ArtCraft anyway. let semantic_media_files = parse_semantic_media_files(request, &creds, &api_host).await?; let reference_tokens = sources_to_artcraft_tokens( - request.reference_media_sources(), ArtcraftMediaKind::Image, Some(&creds), &api_host, + request.reference_media_sources(), ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger, ).await?; let image_inputs = collect_artcraft_image_inputs(&semantic_media_files, reference_tokens); @@ -215,6 +218,7 @@ async fn enqueue_via_fal( request: &TauriGenerateImageRequest, api_key: &str, app_data_root: &AppDataRoot, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -225,7 +229,7 @@ async fn enqueue_via_fal( )), )?; - let image_inputs = resolve_fal_image_inputs(request, &ApiHost::Storyteller, app_data_root).await?; + let image_inputs = resolve_fal_image_inputs(request, &ApiHost::Storyteller, app_data_root, ledger).await?; let router_request = GenerateImageRequestBuilder { model: router_model, @@ -259,6 +263,7 @@ async fn resolve_fal_image_inputs( request: &TauriGenerateImageRequest, api_host: &ApiHost, app_data_root: &AppDataRoot, + ledger: &AssetUploadLedger, ) -> Result, GenerateError> { let mut canvas_and_scene_tokens: Vec = Vec::new(); @@ -278,7 +283,7 @@ async fn resolve_fal_image_inputs( let maybe_creds = maybe_artcraft_credential.as_ref() .and_then(|credential| storyteller_creds_from_credential(credential).ok()); let reference_urls = image_sources_to_fal_urls( - request.reference_media_sources(), maybe_creds.as_ref(), api_host, + request.reference_media_sources(), maybe_creds.as_ref(), api_host, ledger, ).await?; urls.extend(reference_urls.unwrap_or_default()); diff --git a/crates/artcraftx/src/commands/generate/generate_image/generate_image_command.rs b/crates/artcraftx/src/commands/generate/generate_image/generate_image_command.rs index 6cb9d3ac..1418e401 100644 --- a/crates/artcraftx/src/commands/generate/generate_image/generate_image_command.rs +++ b/crates/artcraftx/src/commands/generate/generate_image/generate_image_command.rs @@ -13,6 +13,8 @@ use crate::events::basic_sendable_event_trait::BasicSendableEvent; use crate::events::functional_events::credits_balance_changed_event::CreditsBalanceChangedEvent; use crate::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent; use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::database::local_files_database::LocalFilesDatabase; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::state::database::task_database::TaskDatabase; use log::{error, info}; use tauri::{AppHandle, State}; @@ -25,18 +27,21 @@ pub async fn generate_image_command( task_database: State<'_, TaskDatabase>, midjourney_live_session: State<'_, MidjourneyLiveSession>, grok_websockets: State<'_, GrokWebsockets>, + local_files_database: State<'_, LocalFilesDatabase>, ) -> Response { info!("generate_image_command called, request: {:?}", request); // Generation is credential-driven: the request names a stored credential, // and we invoke the router for that credential's service. + let ledger = AssetUploadLedger::new(local_files_database.inner().clone()); let result = enqueue_image_generation( &request, &app, &app_data_root, &midjourney_live_session, &grok_websockets, + &ledger, ).await; match result { diff --git a/crates/artcraftx/src/commands/generate/generate_image/higgsfield/enqueue_via_higgsfield.rs b/crates/artcraftx/src/commands/generate/generate_image/higgsfield/enqueue_via_higgsfield.rs index c16f2917..cd8c18d9 100644 --- a/crates/artcraftx/src/commands/generate/generate_image/higgsfield/enqueue_via_higgsfield.rs +++ b/crates/artcraftx/src/commands/generate/generate_image/higgsfield/enqueue_via_higgsfield.rs @@ -23,6 +23,7 @@ use crate::commands::generate::task_enqueue_success::TaskEnqueueSuccess; use crate::commands::utils::api_adapters::models::image::tauri_image_model_to_generation_model::tauri_image_model_to_generation_model; use crate::commands::utils::api_adapters::models::image::tauri_image_model_to_router_model::tauri_image_model_to_router_model; use crate::credentials::auth_credential::AuthCredential; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; /// Enqueue via the router's first-party Higgsfield provider. /// @@ -35,6 +36,7 @@ use crate::credentials::auth_credential::AuthCredential; pub async fn enqueue_via_higgsfield( request: &TauriGenerateImageRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -69,7 +71,9 @@ pub async fn enqueue_via_higgsfield( }; let client = higgsfield_router_client(credential)?; - let response = send_higgsfield_image_request(router_request, &client, &media_url_map).await?; + // References this account already uploaded are reused by media id. + let upload_cache = ledger.for_higgsfield_credential(&credential.id); + let response = send_higgsfield_image_request(router_request, &client, &media_url_map, &upload_cache).await?; let payload = response .get_higgsfield_payload() @@ -121,9 +125,15 @@ mod live_higgsfield_image_tests { use crate::commands::generate::generate_image::tauri_image_model::TauriImageModel; use crate::state::data_dir::app_data_root::AppDataRoot; + use crate::state::database::local_files_database::LocalFilesDatabase; use super::*; + async fn ledger() -> AssetUploadLedger { + let app_data_root = AppDataRoot::create_default().expect("app data root"); + AssetUploadLedger::new(LocalFilesDatabase::connect(&app_data_root).await.expect("local files db")) + } + /// A production ArtCraft media file (see `test_data::web::image_media_tokens`). const JUNO_AT_LAKE_MEDIA_TOKEN: &str = "m_m1bz02z1kkzanxy6rb4vk1kvq9de9g"; @@ -152,7 +162,7 @@ mod live_higgsfield_image_tests { ..Default::default() }; - let success = enqueue_via_higgsfield(&request, &credential).await.expect("enqueue should succeed"); + let success = enqueue_via_higgsfield(&request, &credential, &ledger().await).await.expect("enqueue should succeed"); println!("[live] Higgsfield image enqueued: provider_job_id={:?}", success.provider_job_id); assert_eq!(success.provider, GenerationSource::Higgsfield); assert!(success.provider_job_id.is_some()); @@ -171,7 +181,7 @@ mod live_higgsfield_image_tests { ..Default::default() }; - let success = enqueue_via_higgsfield(&request, &credential).await.expect("enqueue should succeed"); + let success = enqueue_via_higgsfield(&request, &credential, &ledger().await).await.expect("enqueue should succeed"); println!("[live] Higgsfield image edit enqueued: provider_job_id={:?}", success.provider_job_id); assert!(success.provider_job_id.is_some()); } diff --git a/crates/artcraftx/src/commands/generate/generate_mesh/generate_mesh_command.rs b/crates/artcraftx/src/commands/generate/generate_mesh/generate_mesh_command.rs index aaa87374..dea54bec 100644 --- a/crates/artcraftx/src/commands/generate/generate_mesh/generate_mesh_command.rs +++ b/crates/artcraftx/src/commands/generate/generate_mesh/generate_mesh_command.rs @@ -13,6 +13,8 @@ use crate::events::generation_events::generation_enqueue_success_event::Generati use crate::state::usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; use crate::state::usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::database::local_files_database::LocalFilesDatabase; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::state::database::task_database::TaskDatabase; use log::{error, info, warn}; use tauri::{AppHandle, State}; @@ -24,15 +26,18 @@ pub async fn generate_mesh_command( app_data_root: State<'_, AppDataRoot>, artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, task_database: State<'_, TaskDatabase>, + local_files_database: State<'_, LocalFilesDatabase>, ) -> Response { info!("generate_mesh_command called, request: {:?}", request); + let ledger = AssetUploadLedger::new(local_files_database.inner().clone()); let result = handle_request( request, &app_data_root, &artcraft_usage_tracker, &task_database, + &ledger, ).await; match result { @@ -92,11 +97,12 @@ async fn handle_request( app_data_root: &AppDataRoot, artcraft_usage_tracker: &ArtcraftUsageTracker, task_database: &TaskDatabase, + ledger: &AssetUploadLedger, ) -> Result { // Generation is credential-driven: the request names a stored credential, // and the router dispatches to that credential's service. - let success_event = handle_credential_router(&request, app_data_root).await?; + let success_event = handle_credential_router(&request, app_data_root, ledger).await?; let result = success_event .insert_into_task_database_with_frontend_payload( diff --git a/crates/artcraftx/src/commands/generate/generate_mesh/handle_credential_router.rs b/crates/artcraftx/src/commands/generate/generate_mesh/handle_credential_router.rs index ca3551bd..c95b084c 100644 --- a/crates/artcraftx/src/commands/generate/generate_mesh/handle_credential_router.rs +++ b/crates/artcraftx/src/commands/generate/generate_mesh/handle_credential_router.rs @@ -14,6 +14,7 @@ use router::api::mesh_ref::MeshRef; use crate::commands::utils::api_adapters::models::mesh::tauri_mesh_model_to_generation_model::tauri_mesh_model_to_generation_model; use crate::commands::utils::api_adapters::models::mesh::tauri_mesh_model_to_router_model::tauri_mesh_model_to_router_model; use crate::commands::generate::generate_error::GenerateError; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::commands::generate::task_enqueue_success::TaskEnqueueSuccess; use crate::commands::generate::common::generation_credential::{credential_not_usable, resolve_generation_credential, storyteller_creds_from_credential}; use crate::commands::generate::common::media_source_conversion::{ @@ -31,6 +32,7 @@ use crate::state::data_dir::app_data_root::AppDataRoot; pub async fn handle_credential_router( request: &TauriGenerateMeshRequest, app_data_root: &AppDataRoot, + ledger: &AssetUploadLedger, ) -> Result { // Reject unusable media sources (missing local files, empty bytes) before // any provider work. @@ -51,7 +53,7 @@ pub async fn handle_credential_router( GenerationSource::Artcraft | GenerationSource::ArtcraftLocal | GenerationSource::ArtcraftCookies => { - handle_artcraft_credential(request, &credential).await + handle_artcraft_credential(request, &credential, ledger).await } other => Err(credential_not_usable( &credential, @@ -66,6 +68,7 @@ pub async fn handle_credential_router( async fn handle_artcraft_credential( request: &TauriGenerateMeshRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -78,15 +81,15 @@ async fn handle_artcraft_credential( let creds = storyteller_creds_from_credential(credential)?; let media = request.media_sources(); - let reference_images = sources_to_artcraft_tokens(media.reference_images, ArtcraftMediaKind::Image, Some(&creds), &api_host) + let reference_images = sources_to_artcraft_tokens(media.reference_images, ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger) .await?.map(ImageListRef::MediaFileTokens); - let front_image = maybe_source_to_artcraft_token(media.front_image, ArtcraftMediaKind::Image, Some(&creds), &api_host) + let front_image = maybe_source_to_artcraft_token(media.front_image, ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger) .await?.map(ImageRef::MediaFileToken); - let back_image = maybe_source_to_artcraft_token(media.back_image, ArtcraftMediaKind::Image, Some(&creds), &api_host) + let back_image = maybe_source_to_artcraft_token(media.back_image, ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger) .await?.map(ImageRef::MediaFileToken); - let left_image = maybe_source_to_artcraft_token(media.left_image, ArtcraftMediaKind::Image, Some(&creds), &api_host) + let left_image = maybe_source_to_artcraft_token(media.left_image, ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger) .await?.map(ImageRef::MediaFileToken); - let right_image = maybe_source_to_artcraft_token(media.right_image, ArtcraftMediaKind::Image, Some(&creds), &api_host) + let right_image = maybe_source_to_artcraft_token(media.right_image, ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger) .await?.map(ImageRef::MediaFileToken); // There's no ArtCraft mesh-file upload endpoint yet; meshes stay token-only. let input_mesh = match media.input_mesh { diff --git a/crates/artcraftx/src/commands/generate/generate_splat/generate_splat_command.rs b/crates/artcraftx/src/commands/generate/generate_splat/generate_splat_command.rs index 1ea15ae1..26074b4d 100644 --- a/crates/artcraftx/src/commands/generate/generate_splat/generate_splat_command.rs +++ b/crates/artcraftx/src/commands/generate/generate_splat/generate_splat_command.rs @@ -13,6 +13,8 @@ use crate::events::generation_events::generation_enqueue_success_event::Generati use crate::state::usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; use crate::state::usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::database::local_files_database::LocalFilesDatabase; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::state::database::task_database::TaskDatabase; use log::{error, info, warn}; use tauri::{AppHandle, State}; @@ -24,15 +26,18 @@ pub async fn generate_splat_command( app_data_root: State<'_, AppDataRoot>, artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, task_database: State<'_, TaskDatabase>, + local_files_database: State<'_, LocalFilesDatabase>, ) -> Response { info!("generate_splat_command called, request: {:?}", request); + let ledger = AssetUploadLedger::new(local_files_database.inner().clone()); let result = handle_request( request, &app_data_root, &artcraft_usage_tracker, &task_database, + &ledger, ).await; match result { @@ -92,11 +97,12 @@ async fn handle_request( app_data_root: &AppDataRoot, artcraft_usage_tracker: &ArtcraftUsageTracker, task_database: &TaskDatabase, + ledger: &AssetUploadLedger, ) -> Result { // Generation is credential-driven: the request names a stored credential, // and the router dispatches to that credential's service. - let success_event = handle_credential_router(&request, app_data_root).await?; + let success_event = handle_credential_router(&request, app_data_root, ledger).await?; let result = success_event .insert_into_task_database_with_frontend_payload( diff --git a/crates/artcraftx/src/commands/generate/generate_splat/handle_credential_router.rs b/crates/artcraftx/src/commands/generate/generate_splat/handle_credential_router.rs index c1500948..dd600f83 100644 --- a/crates/artcraftx/src/commands/generate/generate_splat/handle_credential_router.rs +++ b/crates/artcraftx/src/commands/generate/generate_splat/handle_credential_router.rs @@ -13,6 +13,7 @@ use router::api::video_ref::VideoRef; use crate::commands::utils::api_adapters::models::splat::tauri_splat_model_to_generation_model::tauri_splat_model_to_generation_model; use crate::commands::utils::api_adapters::models::splat::tauri_splat_model_to_router_model::tauri_splat_model_to_router_model; use crate::commands::generate::generate_error::GenerateError; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::commands::generate::task_enqueue_success::TaskEnqueueSuccess; use crate::commands::generate::common::generation_credential::{credential_not_usable, resolve_generation_credential, storyteller_creds_from_credential}; use crate::commands::generate::common::media_source_conversion::{ @@ -30,6 +31,7 @@ use crate::state::data_dir::app_data_root::AppDataRoot; pub async fn handle_credential_router( request: &TauriGenerateSplatRequest, app_data_root: &AppDataRoot, + ledger: &AssetUploadLedger, ) -> Result { // Reject unusable media sources (missing local files, empty bytes) before // any provider work. @@ -50,7 +52,7 @@ pub async fn handle_credential_router( GenerationSource::Artcraft | GenerationSource::ArtcraftLocal | GenerationSource::ArtcraftCookies => { - handle_artcraft_credential(request, &credential).await + handle_artcraft_credential(request, &credential, ledger).await } other => Err(credential_not_usable( &credential, @@ -65,6 +67,7 @@ pub async fn handle_credential_router( async fn handle_artcraft_credential( request: &TauriGenerateSplatRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -77,9 +80,9 @@ async fn handle_artcraft_credential( let creds = storyteller_creds_from_credential(credential)?; let media = request.media_sources(); - let reference_images = sources_to_artcraft_tokens(media.reference_images, ArtcraftMediaKind::Image, Some(&creds), &api_host) + let reference_images = sources_to_artcraft_tokens(media.reference_images, ArtcraftMediaKind::Image, Some(&creds), &api_host, ledger) .await?.map(ImageListRef::MediaFileTokens); - let reference_video = maybe_source_to_artcraft_token(media.reference_video, ArtcraftMediaKind::Video, Some(&creds), &api_host) + let reference_video = maybe_source_to_artcraft_token(media.reference_video, ArtcraftMediaKind::Video, Some(&creds), &api_host, ledger) .await?.map(VideoRef::MediaFileToken); let client = RouterClient::Artcraft(RouterArtcraftClient::new( diff --git a/crates/artcraftx/src/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs b/crates/artcraftx/src/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs index 06c376a4..37b0c8d9 100644 --- a/crates/artcraftx/src/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs +++ b/crates/artcraftx/src/commands/generate/generate_video/artcraft/handle_artcraft_video_via_router.rs @@ -2,6 +2,7 @@ use crate::commands::generate::common::media_source_conversion::{ maybe_source_to_artcraft_token, sources_to_artcraft_tokens, ArtcraftMediaKind, }; use crate::commands::generate::generate_error::GenerateError; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::commands::generate::generate_video::request::TauriGenerateVideoRequest; use crate::commands::generate::task_enqueue_success::TaskEnqueueSuccess; use crate::events::generation_events::common::GenerationModel; @@ -30,6 +31,7 @@ pub async fn handle_artcraft_video_via_router( creds: &StorytellerCredentialSet, model: RouterVideoModel, generation_model: GenerationModel, + ledger: &AssetUploadLedger, ) -> Result { let client = RouterClient::Artcraft(RouterArtcraftClient::new( api_host.clone(), @@ -39,16 +41,16 @@ pub async fn handle_artcraft_video_via_router( // ArtCraft's API is token-native: local files and bytes upload to // ArtCraft here, at generate time (the media must reach ArtCraft anyway). let media = request.media_sources(); - let start_frame = maybe_source_to_artcraft_token(media.start_frame, ArtcraftMediaKind::Image, Some(creds), api_host) + let start_frame = maybe_source_to_artcraft_token(media.start_frame, ArtcraftMediaKind::Image, Some(creds), api_host, ledger) .await?.map(ImageRef::MediaFileToken); - let end_frame = maybe_source_to_artcraft_token(media.end_frame, ArtcraftMediaKind::Image, Some(creds), api_host) + let end_frame = maybe_source_to_artcraft_token(media.end_frame, ArtcraftMediaKind::Image, Some(creds), api_host, ledger) .await?.map(ImageRef::MediaFileToken); - let reference_images = sources_to_artcraft_tokens(media.reference_images, ArtcraftMediaKind::Image, Some(creds), api_host) + let reference_images = sources_to_artcraft_tokens(media.reference_images, ArtcraftMediaKind::Image, Some(creds), api_host, ledger) .await?.map(ImageListRef::MediaFileTokens); - let reference_videos = sources_to_artcraft_tokens(media.reference_videos, ArtcraftMediaKind::Video, Some(creds), api_host) + let reference_videos = sources_to_artcraft_tokens(media.reference_videos, ArtcraftMediaKind::Video, Some(creds), api_host, ledger) .await?.map(VideoListRef::MediaFileTokens); - let reference_audio = sources_to_artcraft_tokens(media.reference_audios, ArtcraftMediaKind::Audio, Some(creds), api_host) + let reference_audio = sources_to_artcraft_tokens(media.reference_audios, ArtcraftMediaKind::Audio, Some(creds), api_host, ledger) .await?.map(AudioListRef::MediaFileTokens); let reference_character_tokens = request.reference_character_tokens.clone().map(CharacterListRef::CharacterTokens); diff --git a/crates/artcraftx/src/commands/generate/generate_video/generate_video_command.rs b/crates/artcraftx/src/commands/generate/generate_video/generate_video_command.rs index 4699a2c6..00f9d32b 100644 --- a/crates/artcraftx/src/commands/generate/generate_video/generate_video_command.rs +++ b/crates/artcraftx/src/commands/generate/generate_video/generate_video_command.rs @@ -14,6 +14,8 @@ use crate::events::generation_events::generation_enqueue_success_event::Generati use crate::state::usage_tracker::artcraft_usage_tracker::ArtcraftUsageTracker; use crate::state::usage_tracker::artcraft_usage_type::{ArtcraftUsagePage, ArtcraftUsageType}; use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::state::database::local_files_database::LocalFilesDatabase; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::state::database::task_database::TaskDatabase; use log::{error, info, warn}; use tauri::{AppHandle, State}; @@ -25,6 +27,7 @@ pub async fn generate_video_command( app_data_root: State<'_, AppDataRoot>, artcraft_usage_tracker: State<'_, ArtcraftUsageTracker>, task_database: State<'_, TaskDatabase>, + local_files_database: State<'_, LocalFilesDatabase>, ) -> Response { info!("generate_video_command called, request: {:?}", request); @@ -41,12 +44,14 @@ pub async fn generate_video_command( request.start_frame_image_media_token = request.image_media_token.clone(); } + let ledger = AssetUploadLedger::new(local_files_database.inner().clone()); let result = handle_request( request, &app, &app_data_root, &artcraft_usage_tracker, &task_database, + &ledger, ).await; match result { @@ -126,6 +131,7 @@ async fn handle_request( app_data_root: &AppDataRoot, artcraft_usage_tracker: &ArtcraftUsageTracker, task_database: &TaskDatabase, + ledger: &AssetUploadLedger, ) -> Result { // Reject unusable media sources (missing local files, empty bytes) before @@ -134,7 +140,7 @@ async fn handle_request( // Generation is credential-driven: the request names a stored credential, // and the router dispatches to that credential's service. - let result = handle_credential_router(&request, Some(app), app_data_root).await; + let result = handle_credential_router(&request, Some(app), app_data_root, ledger).await; let success_event = match result { Err(err) => return Err(err), diff --git a/crates/artcraftx/src/commands/generate/generate_video/handle_credential_router.rs b/crates/artcraftx/src/commands/generate/generate_video/handle_credential_router.rs index 6130866c..0b20074b 100644 --- a/crates/artcraftx/src/commands/generate/generate_video/handle_credential_router.rs +++ b/crates/artcraftx/src/commands/generate/generate_video/handle_credential_router.rs @@ -12,6 +12,7 @@ use crate::utils::services::artcraft_api_host::maybe_artcraft_api_host_for_servi use crate::credentials::auth_credential::AuthCredential; use core_types::enums::generation_source::GenerationSource; use crate::events::notice_events::progress_notice_event::ProgressNotice; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::state::data_dir::app_data_root::AppDataRoot; use tauri::AppHandle; @@ -23,6 +24,7 @@ pub async fn handle_credential_router( // Optional so headless callers (live tests) can skip UI notices. maybe_app: Option<&AppHandle>, app_data_root: &AppDataRoot, + ledger: &AssetUploadLedger, ) -> Result { let credential = resolve_generation_credential( request.credential_id.as_deref(), @@ -39,7 +41,7 @@ pub async fn handle_credential_router( GenerationSource::Artcraft | GenerationSource::ArtcraftLocal | GenerationSource::ArtcraftCookies => { - handle_artcraft_credential(request, &credential).await + handle_artcraft_credential(request, &credential, ledger).await } GenerationSource::HiggsfieldCookies | GenerationSource::Higgsfield => { @@ -51,7 +53,7 @@ pub async fn handle_credential_router( let _ip_check_notice = maybe_app.filter(|_| has_media).map(|app| { ProgressNotice::start(app, "Higgsfield is checking your media for Intellectual Property and Likeness") }); - handle_higgsfield_video_via_router(request, &credential).await + handle_higgsfield_video_via_router(request, &credential, ledger).await } other => Err(credential_not_usable( &credential, @@ -66,6 +68,7 @@ pub async fn handle_credential_router( async fn handle_artcraft_credential( request: &TauriGenerateVideoRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; @@ -83,6 +86,7 @@ async fn handle_artcraft_credential( &creds, router_model, generation_model, + ledger, ).await } @@ -153,7 +157,8 @@ mod live_generation_tests { ..Default::default() }; - let result = handle_credential_router(&request, None, &app_data_root).await; + let ledger = AssetUploadLedger::new(crate::state::database::local_files_database::LocalFilesDatabase::connect(&app_data_root).await.expect("local files db")); + let result = handle_credential_router(&request, None, &app_data_root, &ledger).await; let success = result.expect("seedance lite enqueue should succeed"); println!("[live] seedance 1.0 lite enqueued: job_id={:?}", success.provider_job_id); assert!(success.provider_job_id.is_some()); diff --git a/crates/artcraftx/src/commands/generate/generate_video/higgsfield/handle_higgsfield_video_via_router.rs b/crates/artcraftx/src/commands/generate/generate_video/higgsfield/handle_higgsfield_video_via_router.rs index ad8e898f..2c2fae03 100644 --- a/crates/artcraftx/src/commands/generate/generate_video/higgsfield/handle_higgsfield_video_via_router.rs +++ b/crates/artcraftx/src/commands/generate/generate_video/higgsfield/handle_higgsfield_video_via_router.rs @@ -20,6 +20,7 @@ use crate::commands::generate::task_enqueue_success::TaskEnqueueSuccess; use crate::commands::utils::api_adapters::models::video::tauri_video_model_to_generation_model::tauri_video_model_to_generation_model; use crate::commands::utils::api_adapters::models::video::tauri_video_model_to_router_model::tauri_video_model_to_router_model; use crate::credentials::auth_credential::AuthCredential; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; /// Enqueue via the router's first-party Higgsfield provider. /// @@ -33,6 +34,7 @@ use crate::credentials::auth_credential::AuthCredential; pub async fn handle_higgsfield_video_via_router( request: &TauriGenerateVideoRequest, credential: &AuthCredential, + ledger: &AssetUploadLedger, ) -> Result { let tauri_model = request.model.ok_or(GenerateError::no_model_specified())?; let router_model = tauri_video_model_to_router_model(tauri_model); @@ -65,7 +67,9 @@ pub async fn handle_higgsfield_video_via_router( }; let client = higgsfield_router_client(credential)?; - let response = send_higgsfield_video_request(router_request, &client, &media_url_map).await?; + // References this account already uploaded are reused by media id. + let upload_cache = ledger.for_higgsfield_credential(&credential.id); + let response = send_higgsfield_video_request(router_request, &client, &media_url_map, &upload_cache).await?; let payload = response .get_higgsfield_payload() @@ -104,6 +108,11 @@ mod live_higgsfield_video_tests { use super::*; + async fn test_ledger() -> AssetUploadLedger { + let app_data_root = AppDataRoot::create_default().expect("app data root"); + AssetUploadLedger::new(crate::state::database::local_files_database::LocalFilesDatabase::connect(&app_data_root).await.expect("local files db")) + } + /// A production ArtCraft media file (see `test_data::web::image_media_tokens`). const JUNO_AT_LAKE_MEDIA_TOKEN: &str = "m_m1bz02z1kkzanxy6rb4vk1kvq9de9g"; @@ -133,7 +142,7 @@ mod live_higgsfield_video_tests { ..Default::default() }; - let success = handle_higgsfield_video_via_router(&request, &credential).await.expect("enqueue should succeed"); + let success = handle_higgsfield_video_via_router(&request, &credential, &test_ledger().await).await.expect("enqueue should succeed"); println!("[live] Higgsfield video enqueued: provider_job_id={:?}", success.provider_job_id); assert_eq!(success.provider, GenerationSource::Higgsfield); assert!(success.provider_job_id.is_some()); @@ -153,7 +162,7 @@ mod live_higgsfield_video_tests { ..Default::default() }; - let success = handle_higgsfield_video_via_router(&request, &credential).await.expect("enqueue should succeed"); + let success = handle_higgsfield_video_via_router(&request, &credential, &test_ledger().await).await.expect("enqueue should succeed"); println!("[live] Higgsfield image-to-video enqueued: provider_job_id={:?}", success.provider_job_id); assert!(success.provider_job_id.is_some()); } diff --git a/crates/artcraftx/src/commands/service/promptbox/update_promptbox_state_command.rs b/crates/artcraftx/src/commands/service/promptbox/update_promptbox_state_command.rs index 9333b009..0e8fc520 100644 --- a/crates/artcraftx/src/commands/service/promptbox/update_promptbox_state_command.rs +++ b/crates/artcraftx/src/commands/service/promptbox/update_promptbox_state_command.rs @@ -1,49 +1,27 @@ -use crate::state::promptbox::promptbox_modality::PromptboxModality; use crate::state::promptbox::promptbox_state::PromptboxState; use crate::state::promptbox::promptbox_state_manager::PromptboxStateManager; use log::{error, info}; -use serde_derive::{Deserialize, Serialize}; -use serde_json::{Map, Value}; -use std::collections::BTreeMap; +use serde_derive::Serialize; use tauri::State; -/// A patch to the prompt box state. Every field is optional and only the -/// present ones change; the frontend sends whatever just changed. -#[derive(Deserialize, Debug, Default)] -#[serde(default)] -pub struct UpdatePromptboxStateRequest { - /// Which prompt box the account / model / options apply to. Required when - /// any of those three are present. - pub modality: Option, - pub selected_account_id: Option, - pub selected_model: Option, - /// Replaces the modality's options wholesale (the frontend owns the shape). - pub options: Option>, - /// Merged into `last_account_by_model` (model id -> credential id). - pub last_account_by_model: Option>, -} - #[derive(Serialize)] pub struct UpdatePromptboxStateResponse { pub state: PromptboxState, } +/// Replaces the whole prompt box state. The frontend is the only writer and +/// always knows the full state, so it sends a complete snapshot (debounced) +/// rather than patches. Missing fields default; `version` is owned by the +/// backend and ignored. #[tauri::command] pub async fn update_promptbox_state_command( - request: UpdatePromptboxStateRequest, + state: PromptboxState, promptbox_state: State<'_, PromptboxStateManager>, ) -> Result { - info!("update_promptbox_state_command called: {:?}", request); - - let has_modality_fields = request.selected_account_id.is_some() - || request.selected_model.is_some() - || request.options.is_some(); - if has_modality_fields && request.modality.is_none() { - return Err("`modality` is required when updating account, model, or options".to_string()); - } + info!("update_promptbox_state_command called: {:?}", state); let state = promptbox_state - .update(|state| apply(request, state)) + .update(|current| *current = replace_keeping_version(current, state)) .map_err(|err| { error!("Error updating prompt box state: {:?}", err); format!("Error updating prompt box state: {:?}", err) @@ -52,61 +30,29 @@ pub async fn update_promptbox_state_command( Ok(UpdatePromptboxStateResponse { state }) } -fn apply(request: UpdatePromptboxStateRequest, state: &mut PromptboxState) { - if let Some(modality) = request.modality { - let target = state.modality_mut(modality); - if let Some(account_id) = request.selected_account_id { - target.selected_account_id = Some(account_id); - } - if let Some(model) = request.selected_model { - target.selected_model = Some(model); - } - if let Some(options) = request.options { - target.options = options; - } - } - if let Some(entries) = request.last_account_by_model { - state.last_account_by_model.extend(entries); - } +fn replace_keeping_version(current: &PromptboxState, incoming: PromptboxState) -> PromptboxState { + PromptboxState { version: current.version, ..incoming } } #[cfg(test)] mod tests { use super::*; - use serde_json::json; - - fn decode(json: &str) -> UpdatePromptboxStateRequest { - serde_json::from_str(json).unwrap() - } #[test] - fn patch_only_touches_present_fields() { - let mut state = PromptboxState::default(); - state.image.selected_account_id = Some("credential_a".to_string()); - state.image.options = json!({ "batch_count": 2 }).as_object().unwrap().clone(); + fn incoming_state_replaces_everything_but_version() { + let mut current = PromptboxState::default(); + current.image.selected_model = Some("old_model".to_string()); + current.last_account_by_model.insert("old_model".to_string(), "credential_a".to_string()); - apply(decode(r#"{ "modality": "image", "selected_model": "midjourney_7" }"#), &mut state); - assert_eq!(state.image.selected_model.as_deref(), Some("midjourney_7")); - assert_eq!(state.image.selected_account_id.as_deref(), Some("credential_a"), "untouched"); - assert_eq!(state.image.options["batch_count"], 2, "untouched"); + let incoming: PromptboxState = serde_json::from_str( + r#"{ "version": 999, "video": { "selected_model": "seedance_2p0", "options": { "duration": 5 } } }"#, + ).unwrap(); - apply(decode(r#"{ "modality": "image", "options": { "aspect_ratio": "square" } }"#), &mut state); - assert_eq!(state.image.options.len(), 1, "options replace wholesale"); - assert_eq!(state.image.options["aspect_ratio"], "square"); - assert_eq!(state.video, Default::default(), "other modalities untouched"); - } - - #[test] - fn last_account_entries_merge() { - let mut state = PromptboxState::default(); - apply(decode(r#"{ "last_account_by_model": { "midjourney_7": "credential_a" } }"#), &mut state); - apply(decode(r#"{ "last_account_by_model": { "grok_imagine_image": "credential_g", "midjourney_7": "credential_b" } }"#), &mut state); - assert_eq!(state.last_account_by_model["midjourney_7"], "credential_b"); - assert_eq!(state.last_account_by_model["grok_imagine_image"], "credential_g"); - } - - #[test] - fn unknown_modality_is_rejected() { - assert!(serde_json::from_str::(r#"{ "modality": "hologram" }"#).is_err()); + let replaced = replace_keeping_version(¤t, incoming); + assert_eq!(replaced.version, current.version, "backend owns the version"); + assert_eq!(replaced.video.selected_model.as_deref(), Some("seedance_2p0")); + assert_eq!(replaced.video.options["duration"], 5); + assert_eq!(replaced.image, Default::default(), "omitted modalities reset, not merged"); + assert!(replaced.last_account_by_model.is_empty(), "omitted map resets, not merged"); } } diff --git a/crates/artcraftx/src/services/asset_uploads/asset_upload_ledger.rs b/crates/artcraftx/src/services/asset_uploads/asset_upload_ledger.rs new file mode 100644 index 00000000..8b1a700e --- /dev/null +++ b/crates/artcraftx/src/services/asset_uploads/asset_upload_ledger.rs @@ -0,0 +1,316 @@ +//! The upload-dedup ledger: which remote service already holds which local +//! content, per account. +//! +//! Before an upload, callers hash the bytes (memoized through the +//! local_files index for files on disk) and ask the ledger. A hit is +//! verified with the service before it is trusted; a miss, or a hit the +//! service no longer honours, means a fresh upload that is then recorded +//! over the old row. The ledger never makes an upload fail: every lookup +//! error reads as a miss and every write error is logged. +//! +//! Account identity: +//! - Higgsfield: the stored credential's id (sessions only come from stored +//! credentials). +//! - ArtCraft: a fingerprint of the web session cookie +//! ([`artcraft_account_id`]) — the app has no user id for a session, and a +//! re-login simply starts a fresh (empty) history. +//! +//! The router gets a per-account view through [`AccountUploadCache`], which +//! implements its [`AssetUploadCache`] trait. + +use std::future::Future; +use std::path::Path; + +use artcraft_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use artcraft_client::endpoints::media_files::get_media_file::get_media_file; +use artcraft_client::utils::api_host::ApiHost; +use async_trait::async_trait; +use core_types::enums::generation_source::GenerationSource; +use core_types::identifiers::credential_id::CredentialId; +use file_hashing::hash_bytes_blake3; +use local_files_database::queries::asset_upload_record::AssetUploadRecord; +use local_files_database::queries::delete_asset_upload::delete_asset_upload; +use local_files_database::queries::get_asset_upload::{get_asset_upload, GetAssetUploadArgs}; +use local_files_database::queries::touch_asset_upload_reused::touch_asset_upload_reused; +use local_files_database::queries::upsert_asset_upload::{upsert_asset_upload, UpsertAssetUploadArgs}; +use log::{info, warn}; +use router::api::asset_upload_cache::{AssetUploadCache, CachedAssetUpload}; +use sqlite_identifiers::ids::media_file_token::MediaFileToken; + +use crate::services::local_files::file_hash_service::hash_local_file_memoized; +use crate::state::database::local_files_database::LocalFilesDatabase; + +/// The ArtCraft account key is derived from the session cookie; this many +/// hex characters of its BLAKE3 is plenty to tell sessions apart. +const ARTCRAFT_SESSION_FINGERPRINT_CHARS: usize = 32; + +#[derive(Clone)] +pub struct AssetUploadLedger { + database: LocalFilesDatabase, +} + +/// The ledger narrowed to one service + account: what the router (and the +/// ArtCraft upload helpers) work with. +#[derive(Clone)] +pub struct AccountUploadCache { + ledger: AssetUploadLedger, + service: GenerationSource, + account_id: String, +} + +impl AssetUploadLedger { + pub fn new(database: LocalFilesDatabase) -> Self { + Self { database } + } + + /// This account's uploads on `service`. + pub fn for_account(&self, service: GenerationSource, account_id: impl Into) -> AccountUploadCache { + AccountUploadCache { ledger: self.clone(), service, account_id: account_id.into() } + } + + /// A stored Higgsfield credential's uploads. + pub fn for_higgsfield_credential(&self, credential_id: &CredentialId) -> AccountUploadCache { + self.for_account(GenerationSource::Higgsfield, credential_id.as_str()) + } + + /// The uploads of the ArtCraft account behind a web session, if the + /// session identifies one (anonymous uploads aren't tracked). + pub fn for_artcraft_session(&self, maybe_creds: Option<&StorytellerCredentialSet>) -> Option { + let account_id = artcraft_account_id(maybe_creds?)?; + Some(self.for_account(GenerationSource::Artcraft, account_id)) + } + + /// BLAKE3 of a file on disk, memoized in the local_files index. + pub async fn hash_file(&self, path: &Path) -> Result { + hash_local_file_memoized(&self.database, path).await + } + + // ── Raw row access ── + + pub async fn find(&self, service: GenerationSource, account_id: &str, file_hash: &str) -> Option { + let result = get_asset_upload(GetAssetUploadArgs { + db: self.database.get_connection(), + service: service.to_str(), + account_id, + file_hash_blake3: file_hash, + }).await; + match result { + Ok(record) => record, + Err(err) => { + warn!("asset_uploads lookup failed ({service}, {account_id}, {file_hash}): {err}"); + None + } + } + } + + pub async fn record(&self, service: GenerationSource, account_id: &str, file_hash: &str, file_size_bytes: Option, upload: &CachedAssetUpload) { + let result = upsert_asset_upload(UpsertAssetUploadArgs { + db: self.database.get_connection(), + service: service.to_str(), + account_id, + file_hash_blake3: file_hash, + service_id: &upload.service_id, + maybe_service_url: upload.maybe_service_url.as_deref(), + maybe_file_size_bytes: file_size_bytes.map(|size| size as i64), + }).await; + match result { + Ok(()) => info!("Recorded {service} upload {} for {account_id} (hash {file_hash})", upload.service_id), + Err(err) => warn!("Could not record {service} upload {}: {err}", upload.service_id), + } + } + + pub async fn note_reused(&self, service: GenerationSource, account_id: &str, file_hash: &str) { + if let Err(err) = touch_asset_upload_reused(self.database.get_connection(), service.to_str(), account_id, file_hash).await { + warn!("Could not mark {service} upload reused (hash {file_hash}): {err}"); + } + } + + pub async fn forget(&self, service: GenerationSource, account_id: &str, file_hash: &str) { + if let Err(err) = delete_asset_upload(self.database.get_connection(), service.to_str(), account_id, file_hash).await { + warn!("Could not forget {service} upload (hash {file_hash}): {err}"); + } + } +} + +impl AccountUploadCache { + pub fn service(&self) -> GenerationSource { + self.service + } + + pub fn account_id(&self) -> &str { + &self.account_id + } + + /// Hash a local file (memoized through the local_files index). + pub async fn hash_file(&self, path: &Path) -> Result { + self.ledger.hash_file(path).await + } + + /// The ArtCraft media file for a local file, reusing this account's + /// earlier upload when ArtCraft still has it, else running `upload` and + /// recording its token. `upload` is only awaited on a miss. + /// + /// Hashing failures (unreadable file) skip the ledger and just upload. + pub async fn artcraft_token_for_file(&self, api_host: &ApiHost, path: &Path, upload: F) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let maybe_hash = match self.hash_file(path).await { + Ok(hash) => Some(hash), + Err(message) => { + warn!("Could not hash {} for the upload ledger ({message}); uploading without it", path.display()); + None + } + }; + + if let Some(hash) = maybe_hash.as_deref() { + if let Some(token) = self.verified_artcraft_token(api_host, hash).await { + info!("Reusing ArtCraft media file {} for {} (already uploaded by this account)", token.as_str(), path.display()); + self.note_reused(hash).await; + return Ok(token); + } + } + + let token = upload().await?; + + if let Some(hash) = maybe_hash.as_deref() { + let file_size_bytes = std::fs::metadata(path).ok().map(|metadata| metadata.len()); + self.record_upload(hash, file_size_bytes.unwrap_or(0), &CachedAssetUpload { + service_id: token.as_str().to_string(), + maybe_service_url: None, + }).await; + } + Ok(token) + } + + /// Same as [`Self::artcraft_token_for_file`] for bytes already in memory. + pub async fn artcraft_token_for_bytes(&self, api_host: &ApiHost, bytes: &[u8], upload: F) -> Result + where + F: FnOnce() -> Fut, + Fut: Future>, + { + let hash = hash_bytes_blake3(bytes); + if let Some(token) = self.verified_artcraft_token(api_host, &hash).await { + info!("Reusing ArtCraft media file {} for {} in-memory bytes (already uploaded by this account)", token.as_str(), bytes.len()); + self.note_reused(&hash).await; + return Ok(token); + } + let token = upload().await?; + self.record_upload(&hash, bytes.len() as u64, &CachedAssetUpload { + service_id: token.as_str().to_string(), + maybe_service_url: None, + }).await; + Ok(token) + } + + /// The cached ArtCraft token for a hash, if ArtCraft still serves the + /// media file. Anything else forgets the row and answers `None`. + async fn verified_artcraft_token(&self, api_host: &ApiHost, file_hash: &str) -> Option { + let cached = self.find_upload(file_hash).await?; + let token = MediaFileToken::new_from_str(&cached.service_id); + match get_media_file(api_host, &token).await { + Ok(_) => Some(token), + Err(err) => { + warn!("Cached ArtCraft media file {} could not be confirmed ({:?}); uploading again", token.as_str(), err); + self.forget_upload(file_hash).await; + None + } + } + } +} + +#[async_trait] +impl AssetUploadCache for AccountUploadCache { + async fn find_upload(&self, file_hash_blake3: &str) -> Option { + self.ledger.find(self.service, &self.account_id, file_hash_blake3).await + .map(|record| CachedAssetUpload { + service_id: record.service_id, + maybe_service_url: record.service_url, + }) + } + + async fn record_upload(&self, file_hash_blake3: &str, file_size_bytes: u64, upload: &CachedAssetUpload) { + self.ledger.record(self.service, &self.account_id, file_hash_blake3, Some(file_size_bytes), upload).await + } + + async fn note_reused(&self, file_hash_blake3: &str) { + self.ledger.note_reused(self.service, &self.account_id, file_hash_blake3).await + } + + async fn forget_upload(&self, file_hash_blake3: &str) { + self.ledger.forget(self.service, &self.account_id, file_hash_blake3).await + } +} + +/// The ledger's account key for an ArtCraft web session: a fingerprint of +/// the session cookie. `None` without a session (anonymous). +pub fn artcraft_account_id(creds: &StorytellerCredentialSet) -> Option { + let session = creds.session.as_ref()?; + let fingerprint = hash_bytes_blake3(session.as_bytes()); + Some(format!("artcraft_session_{}", &fingerprint[..ARTCRAFT_SESSION_FINGERPRINT_CHARS])) +} + +#[cfg(test)] +mod tests { + use super::*; + use artcraft_client::credentials::storyteller_session_cookie::StorytellerSessionCookie; + use crate::state::data_dir::app_data_root::AppDataRoot; + + #[test] + fn artcraft_account_id_is_a_stable_session_fingerprint() { + let mut creds = StorytellerCredentialSet::empty(); + assert!(artcraft_account_id(&creds).is_none()); + + creds.session = Some(StorytellerSessionCookie::new("session-abc".to_string())); + let id = artcraft_account_id(&creds).unwrap(); + assert!(id.starts_with("artcraft_session_")); + assert_eq!(id.len(), "artcraft_session_".len() + ARTCRAFT_SESSION_FINGERPRINT_CHARS); + assert_eq!(artcraft_account_id(&creds).unwrap(), id); + assert!(!id.contains("session-abc"), "the raw cookie must not leak into the key"); + + creds.session = Some(StorytellerSessionCookie::new("session-xyz".to_string())); + assert_ne!(artcraft_account_id(&creds).unwrap(), id); + } + + #[tokio::test] + async fn account_cache_round_trips_through_the_router_trait() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path().join("root")).unwrap(); + let ledger = AssetUploadLedger::new(LocalFilesDatabase::connect(&root).await.unwrap()); + let cache = ledger.for_higgsfield_credential(&CredentialId::from_trusted("credential_test")); + let other = ledger.for_higgsfield_credential(&CredentialId::from_trusted("credential_other")); + + assert!(cache.find_upload("hash1").await.is_none()); + let upload = CachedAssetUpload { service_id: "media-1".into(), maybe_service_url: Some("https://cdn/x".into()) }; + cache.record_upload("hash1", 10, &upload).await; + assert_eq!(cache.find_upload("hash1").await, Some(upload)); + assert!(other.find_upload("hash1").await.is_none(), "accounts don't share uploads"); + + cache.note_reused("hash1").await; + assert!(ledger.find(GenerationSource::Higgsfield, "credential_test", "hash1").await.unwrap().last_reused_at.is_some()); + + cache.forget_upload("hash1").await; + assert!(cache.find_upload("hash1").await.is_none()); + } + + #[tokio::test] + async fn artcraft_file_helper_uploads_on_miss_and_records() { + let dir = tempfile::tempdir().unwrap(); + let root = AppDataRoot::create_existing(dir.path().join("root")).unwrap(); + let ledger = AssetUploadLedger::new(LocalFilesDatabase::connect(&root).await.unwrap()); + let cache = ledger.for_account(GenerationSource::Artcraft, "artcraft_session_test"); + let file = dir.path().join("ref.png"); + std::fs::write(&file, b"\x89PNG\r\n\x1a\nfake").unwrap(); + + let token = cache.artcraft_token_for_file::<_, _, ()>(&ApiHost::Storyteller, &file, || async { + Ok(MediaFileToken::new_from_str("m_fresh")) + }).await.unwrap(); + assert_eq!(token.as_str(), "m_fresh"); + + let hash = cache.hash_file(&file).await.unwrap(); + let recorded = cache.find_upload(&hash).await.unwrap(); + assert_eq!(recorded.service_id, "m_fresh"); + assert_eq!(ledger.find(GenerationSource::Artcraft, "artcraft_session_test", &hash).await.unwrap().file_size_bytes, Some(12)); + } +} diff --git a/crates/artcraftx/src/services/asset_uploads/mod.rs b/crates/artcraftx/src/services/asset_uploads/mod.rs new file mode 100644 index 00000000..719a52be --- /dev/null +++ b/crates/artcraftx/src/services/asset_uploads/mod.rs @@ -0,0 +1 @@ +pub mod asset_upload_ledger; diff --git a/crates/artcraftx/src/services/local_files/file_hash_service.rs b/crates/artcraftx/src/services/local_files/file_hash_service.rs new file mode 100644 index 00000000..e1f2f061 --- /dev/null +++ b/crates/artcraftx/src/services/local_files/file_hash_service.rs @@ -0,0 +1,80 @@ +//! BLAKE3 for local files, memoized in the local_files index: when a path's +//! size and mtime match the indexed row the stored hash is returned without +//! reading the file; otherwise the file is hashed and the row refreshed. + +use std::path::Path; +use std::time::UNIX_EPOCH; + +use file_hashing::hash_file_blake3; +use local_files_database::queries::get_local_file_by_path::get_local_file_by_path; +use local_files_database::queries::upsert_local_file::{upsert_local_file, UpsertLocalFileArgs}; +use log::warn; + +use crate::state::database::local_files_database::LocalFilesDatabase; + +/// `(size, mtime_ms)` for a regular file, `None` if it's missing. +pub fn stat_file(file_path: &Path) -> Option<(i64, i64)> { + let stat = match std::fs::metadata(file_path) { + Ok(metadata) if metadata.is_file() => metadata, + _ => return None, + }; + let file_mtime_ms = stat.modified().ok() + .and_then(|mtime| mtime.duration_since(UNIX_EPOCH).ok()) + .map(|duration| duration.as_millis() as i64) + .unwrap_or(0); + Some((stat.len() as i64, file_mtime_ms)) +} + +/// The file's BLAKE3 hash, from the index when the stat matches, else by +/// reading the file (and indexing the result). +pub async fn hash_local_file_memoized(database: &LocalFilesDatabase, file_path: &Path) -> Result { + let (file_size_bytes, file_mtime_ms) = stat_file(file_path) + .ok_or_else(|| format!("{} is not a readable file", file_path.display()))?; + hash_local_file_with_stat(database, file_path, file_size_bytes, file_mtime_ms).await +} + +/// [`hash_local_file_memoized`] for a caller that already stat-ed the file. +pub async fn hash_local_file_with_stat( + database: &LocalFilesDatabase, + file_path: &Path, + file_size_bytes: i64, + file_mtime_ms: i64, +) -> Result { + let path_string = file_path.to_string_lossy().into_owned(); + + match get_local_file_by_path(database.get_connection(), &path_string).await { + Ok(Some(record)) if record.matches_stat(file_size_bytes, file_mtime_ms) => { + return Ok(record.file_hash_blake3); + } + Ok(_) => {} + Err(err) => warn!("local_files lookup failed for {}: {}", path_string, err), + } + + let hash_path = file_path.to_path_buf(); + let file_hash = tokio::task::spawn_blocking(move || hash_file_blake3(&hash_path)) + .await + .map_err(|join_error| join_error.to_string())? + .map_err(|io_error| io_error.to_string())?; + + let extension = file_extension(file_path); + let upsert = upsert_local_file(UpsertLocalFileArgs { + db: database.get_connection(), + file_path: &path_string, + maybe_file_type: (!extension.is_empty()).then_some(extension.as_str()), + file_size_bytes, + file_mtime_ms, + file_hash_blake3: &file_hash, + }).await; + if let Err(err) = upsert { + warn!("local_files upsert failed for {}: {}", path_string, err); + } + + Ok(file_hash) +} + +pub fn file_extension(file_path: &Path) -> String { + file_path.extension() + .and_then(|extension| extension.to_str()) + .map(str::to_ascii_lowercase) + .unwrap_or_default() +} diff --git a/crates/artcraftx/src/services/local_files/mod.rs b/crates/artcraftx/src/services/local_files/mod.rs index dae4ce99..328e355c 100644 --- a/crates/artcraftx/src/services/local_files/mod.rs +++ b/crates/artcraftx/src/services/local_files/mod.rs @@ -1 +1,2 @@ +pub mod file_hash_service; pub mod thumbnail_service; diff --git a/crates/artcraftx/src/services/local_files/thumbnail_service.rs b/crates/artcraftx/src/services/local_files/thumbnail_service.rs index 6f83eea4..090558e6 100644 --- a/crates/artcraftx/src/services/local_files/thumbnail_service.rs +++ b/crates/artcraftx/src/services/local_files/thumbnail_service.rs @@ -17,13 +17,10 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex as StdMutex}; -use std::time::UNIX_EPOCH; -use file_hashing::hash_file_blake3; use local_files_database::queries::get_local_file_by_path::get_local_file_by_path; use local_files_database::queries::local_file_record::LocalFileRecord; use local_files_database::queries::set_thumbnail_result::{set_thumbnail_result, SetThumbnailResultArgs}; -use local_files_database::queries::upsert_local_file::{upsert_local_file, UpsertLocalFileArgs}; use local_files_database::thumbnail_state::ThumbnailState; use log::{info, warn}; use once_cell::sync::Lazy; @@ -31,6 +28,7 @@ use serde_derive::Serialize; use thumbnails::{generate_image_thumbnail, generate_video_thumbnails, ThumbnailError, STATIC_THUMBNAIL_MAX_DIMENSION, THUMBNAIL_GENERATOR_VERSION}; use tokio::sync::{Mutex as TokioMutex, Semaphore}; +use crate::services::local_files::file_hash_service::{file_extension, hash_local_file_with_stat, stat_file}; use crate::state::data_dir::app_data_root::AppDataRoot; use crate::state::database::local_files_database::LocalFilesDatabase; @@ -184,7 +182,7 @@ pub async fn generate_thumbnail( return LocalThumbnail::without_thumbnail(file_path, LocalThumbnailStatus::MissingFile, None); }; - let file_hash = match cached_or_computed_hash(database, file_path, file_size_bytes, file_mtime_ms).await { + let file_hash = match hash_local_file_with_stat(database, file_path, file_size_bytes, file_mtime_ms).await { Ok(hash) => hash, Err(message) => { warn!("Could not hash {}: {}", file_path.display(), message); @@ -284,47 +282,6 @@ fn generate(kind: MediaKind, source: &Path, jpg: &Path, webp: &Path) -> Result Result { - let path_string = file_path.to_string_lossy().into_owned(); - - match get_local_file_by_path(database.get_connection(), &path_string).await { - Ok(Some(record)) if record.matches_stat(file_size_bytes, file_mtime_ms) => { - return Ok(record.file_hash_blake3); - } - Ok(_) => {} - Err(err) => warn!("local_files lookup failed for {}: {}", path_string, err), - } - - let hash_path = file_path.to_path_buf(); - let file_hash = tokio::task::spawn_blocking(move || hash_file_blake3(&hash_path)) - .await - .map_err(|join_error| join_error.to_string())? - .map_err(|io_error| io_error.to_string())?; - - let extension = file_extension(file_path); - let upsert = upsert_local_file(UpsertLocalFileArgs { - db: database.get_connection(), - file_path: &path_string, - maybe_file_type: (!extension.is_empty()).then_some(extension.as_str()), - file_size_bytes, - file_mtime_ms, - file_hash_blake3: &file_hash, - }).await; - if let Err(err) = upsert { - warn!("local_files upsert failed for {}: {}", path_string, err); - } - - Ok(file_hash) -} - async fn record_outcome( database: &LocalFilesDatabase, file_path: &Path, @@ -352,31 +309,11 @@ async fn record_outcome( } } -/// `(size, mtime_ms)` for a regular file, `None` if it's missing. -fn stat_file(file_path: &Path) -> Option<(i64, i64)> { - let stat = match std::fs::metadata(file_path) { - Ok(metadata) if metadata.is_file() => metadata, - _ => return None, - }; - let file_mtime_ms = stat.modified().ok() - .and_then(|mtime| mtime.duration_since(UNIX_EPOCH).ok()) - .map(|duration| duration.as_millis() as i64) - .unwrap_or(0); - Some((stat.len() as i64, file_mtime_ms)) -} - fn per_path_lock(file_path: &Path) -> Arc> { let mut locks = PER_PATH_LOCKS.lock().expect("lock poisoned"); locks.entry(file_path.to_path_buf()).or_default().clone() } -fn file_extension(file_path: &Path) -> String { - file_path.extension() - .and_then(|extension| extension.to_str()) - .map(str::to_ascii_lowercase) - .unwrap_or_default() -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/artcraftx/src/services/mod.rs b/crates/artcraftx/src/services/mod.rs index 330c9f91..ca61b941 100644 --- a/crates/artcraftx/src/services/mod.rs +++ b/crates/artcraftx/src/services/mod.rs @@ -1,3 +1,4 @@ +pub mod asset_uploads; pub mod grok; pub mod local_files; pub mod higgsfield; diff --git a/crates/artcraftx/src/state/promptbox/promptbox_state_manager.rs b/crates/artcraftx/src/state/promptbox/promptbox_state_manager.rs index 72554075..a89953b6 100644 --- a/crates/artcraftx/src/state/promptbox/promptbox_state_manager.rs +++ b/crates/artcraftx/src/state/promptbox/promptbox_state_manager.rs @@ -1,7 +1,9 @@ use crate::state::data_dir::app_data_root::AppDataRoot; use crate::state::promptbox::promptbox_state::PromptboxState; use errors::AnyhowResult; +use log::warn; use memory_store::clone_cell::CloneCell; +use std::sync::{Arc, Mutex, MutexGuard}; /// Shared, process-wide prompt box state: an in-memory copy that readers /// clone, and a single write path that persists to disk. @@ -9,6 +11,11 @@ use memory_store::clone_cell::CloneCell; pub struct PromptboxStateManager { state: CloneCell, data_root: AppDataRoot, + /// Serializes `update` so concurrent commands don't race each other's + /// read-modify-write or collide on the same temp file when saving. + /// The guarded section is synchronous and tiny (clone, serialize a few KB, + /// write, rename), so callers queue for microseconds, never the UI. + write_lock: Arc>, } impl PromptboxStateManager { @@ -16,6 +23,7 @@ impl PromptboxStateManager { Self { state: CloneCell::with_owned(PromptboxState::load_or_default(data_root)), data_root: data_root.clone(), + write_lock: Arc::new(Mutex::new(())), } } @@ -26,10 +34,62 @@ impl PromptboxStateManager { /// Apply a change, persist it, and return the updated snapshot. Nothing /// changes in memory if the write fails. pub fn update(&self, change: impl FnOnce(&mut PromptboxState)) -> AnyhowResult { + let _guard = self.acquire_write_lock(); let mut state = self.get()?; change(&mut state); state.save(&self.data_root)?; self.state.set_clone(&state)?; Ok(state) } + + /// The lock guards nothing but ordering (the real state lives in the cell + /// and on disk, and saves are atomic), so a poisoned lock is safe to reuse. + /// Failing here forever would just break a low-stakes preference feature. + fn acquire_write_lock(&self) -> MutexGuard<'_, ()> { + self.write_lock.lock().unwrap_or_else(|poisoned| { + warn!("Prompt box write lock was poisoned by an earlier panic; continuing"); + poisoned.into_inner() + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + + const THREADS: usize = 8; + const UPDATES_PER_THREAD: usize = 25; + + /// Regression: concurrent updates used to race on the shared temp file + /// (`rename` -> ENOENT) and clobber each other's changes. + #[test] + fn concurrent_updates_all_succeed_and_none_are_lost() { + let dir = tempfile::tempdir().unwrap(); + let data_root = AppDataRoot::create_existing(dir.path()).unwrap(); + let manager = PromptboxStateManager::load_or_default(&data_root); + + let handles: Vec<_> = (0..THREADS) + .map(|thread_index| { + let manager = manager.clone(); + thread::spawn(move || { + for i in 0..UPDATES_PER_THREAD { + manager + .update(|state| { + state.last_account_by_model.insert(format!("model_{thread_index}_{i}"), "credential".to_string()); + }) + .expect("update should never fail under contention"); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + let expected = THREADS * UPDATES_PER_THREAD; + assert_eq!(manager.get().unwrap().last_account_by_model.len(), expected, "in memory"); + let on_disk = PromptboxState::load_or_default(&data_root); + assert_eq!(on_disk.last_account_by_model.len(), expected, "on disk"); + } } diff --git a/crates/artcraftx/src/threads/task_completion/complete_task_with_local_files.rs b/crates/artcraftx/src/threads/task_completion/complete_task_with_local_files.rs index 74adcd79..37fc8460 100644 --- a/crates/artcraftx/src/threads/task_completion/complete_task_with_local_files.rs +++ b/crates/artcraftx/src/threads/task_completion/complete_task_with_local_files.rs @@ -1,5 +1,7 @@ use crate::state::app_preferences::app_preferences_manager::AppPreferencesManager; use crate::state::data_dir::app_data_root::AppDataRoot; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; +use crate::state::database::local_files_database::LocalFilesDatabase; use crate::state::database::task_database::TaskDatabase; use crate::state::thumbnails::thumbnail_work_queue::ThumbnailWorkQueue; use crate::threads::task_completion::save_results_to_download_dir::save_results_to_download_dir; @@ -20,7 +22,7 @@ use sqlite_database::queries::update::update_successful_task_status_with_metadat }; use sqlite_identifiers::enums::task_media_file_class::TaskMediaFileClass; use std::path::PathBuf; -use tauri::AppHandle; +use tauri::{AppHandle, Manager}; pub struct CompleteTaskArgs<'a> { pub app_handle: &'a AppHandle, @@ -95,6 +97,10 @@ pub async fn complete_task_with_local_files(args: CompleteTaskArgs<'_>) -> Anyho local_files, ); + // The upload ledger lets a re-delivered file reuse its earlier upload. + let maybe_ledger = app_handle.try_state::() + .map(|database| AssetUploadLedger::new(database.inner().clone())); + let maybe_uploaded = match maybe_storyteller_creds { Some(creds) => { let uploaded = upload_results_to_artcraft( @@ -104,6 +110,7 @@ pub async fn complete_task_with_local_files(args: CompleteTaskArgs<'_>) -> Anyho media_class, prompt, local_files, + maybe_ledger.as_ref(), ).await?; Some(uploaded) } diff --git a/crates/artcraftx/src/threads/task_completion/upload_results_to_artcraft.rs b/crates/artcraftx/src/threads/task_completion/upload_results_to_artcraft.rs index dc9991aa..5af3d4b5 100644 --- a/crates/artcraftx/src/threads/task_completion/upload_results_to_artcraft.rs +++ b/crates/artcraftx/src/threads/task_completion/upload_results_to_artcraft.rs @@ -28,6 +28,7 @@ use sqlite_identifiers::ids::prompt_token::PromptToken; use std::path::{Path, PathBuf}; use std::time::Duration; use uuid_utils::uuid::generate_random_uuid; +use crate::services::asset_uploads::asset_upload_ledger::AssetUploadLedger; use crate::utils::enum_conversion::artcraft_api_generation_provider::artcraft_api_generation_provider; const MAX_UPLOAD_RETRIES: u32 = 5; @@ -63,7 +64,8 @@ pub struct UploadedResults { /// Upload every result file to ArtCraft (in order; the first becomes the /// primary), retrying on rate limits, then look up the primary file's CDN and -/// thumbnail URLs (failing open). +/// thumbnail URLs (failing open). With a ledger, a file this account already +/// uploaded (and ArtCraft still has) is reused by token instead. pub async fn upload_results_to_artcraft( creds: &StorytellerCredentialSet, task: &Task, @@ -71,8 +73,10 @@ pub async fn upload_results_to_artcraft( media_class: TaskMediaFileClass, prompt: CompletionPrompt, local_files: &[PathBuf], + maybe_ledger: Option<&AssetUploadLedger>, ) -> AnyhowResult { let task_id = task.id.as_str(); + let maybe_upload_cache = maybe_ledger.and_then(|ledger| ledger.for_artcraft_session(Some(creds))); let maybe_prompt_token = resolve_prompt_token(creds, task, generation_provider, prompt).await; @@ -85,14 +89,18 @@ pub async fn upload_results_to_artcraft( for (index, local_file) in local_files.iter().enumerate() { info!("[TaskCompletion] Uploading result {} of {} for task {} ...", index + 1, local_files.len(), task_id); - let media_token = upload_with_retry( + let upload = || upload_with_retry( creds, local_file, generation_provider, media_class, maybe_prompt_token.as_ref(), maybe_batch_token.as_ref(), - ).await?; + ); + let media_token = match &maybe_upload_cache { + Some(cache) => cache.artcraft_token_for_file(&ApiHost::Storyteller, local_file, upload).await?, + None => upload().await?, + }; info!("[TaskCompletion] Uploaded task {} result as {:?}", task_id, media_token); diff --git a/crates/router/Cargo.toml b/crates/router/Cargo.toml index cd52a5b7..6cc28cac 100644 --- a/crates/router/Cargo.toml +++ b/crates/router/Cargo.toml @@ -17,6 +17,7 @@ workspace-build-acceleration.workspace = true # Internal artcraft_client.workspace = true +file_hashing.workspace = true browser_emulation.workspace = true fal_client.workspace = true gmicloud_client.workspace = true @@ -31,6 +32,7 @@ worldlabs_api_client.workspace = true # External anyhow.workspace = true +async-trait.workspace = true log.workspace = true reqwest.workspace = true serde.workspace = true diff --git a/crates/router/src/api/asset_upload_cache.rs b/crates/router/src/api/asset_upload_cache.rs new file mode 100644 index 00000000..6407c091 --- /dev/null +++ b/crates/router/src/api/asset_upload_cache.rs @@ -0,0 +1,39 @@ +//! Skip re-uploading bytes a provider already holds. +//! +//! The router uploads reference media to providers on the caller's behalf +//! (Higgsfield today). The app keeps a ledger of what each account has +//! already uploaded, keyed by content hash; it hands the router a +//! per-account view of that ledger through this trait, and the router asks +//! before every upload. The router has no idea how the ledger is stored. +//! +//! Entries are advisory: the router verifies a hit with the provider before +//! trusting it, and any failure falls back to a fresh upload that is then +//! recorded over the stale entry. + +use async_trait::async_trait; + +/// A previously uploaded copy of some bytes, as the provider knows it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CachedAssetUpload { + /// The provider's own id for the asset (e.g. a Higgsfield media id). + pub service_id: String, + /// Where the provider serves it, when it publishes a URL. + pub maybe_service_url: Option, +} + +/// One account's upload history on one provider. Implementations must not +/// fail loudly: a lookup error reads as a miss, a record error is logged. +#[async_trait] +pub trait AssetUploadCache: Send + Sync { + /// The upload previously made for these bytes, if any. + async fn find_upload(&self, file_hash_blake3: &str) -> Option; + + /// Remember a fresh upload (replacing any stale entry for the hash). + async fn record_upload(&self, file_hash_blake3: &str, file_size_bytes: u64, upload: &CachedAssetUpload); + + /// Note that an upload was skipped thanks to a (verified) entry. + async fn note_reused(&self, file_hash_blake3: &str); + + /// Drop an entry the provider no longer honours. + async fn forget_upload(&self, file_hash_blake3: &str); +} diff --git a/crates/router/src/api/mod.rs b/crates/router/src/api/mod.rs index c3d1cec3..47a04a9e 100644 --- a/crates/router/src/api/mod.rs +++ b/crates/router/src/api/mod.rs @@ -1,3 +1,4 @@ +pub mod asset_upload_cache; pub mod audio_list_ref; pub mod audio_ref; pub mod character_list_ref; diff --git a/crates/router/src/generate/generate_image/image_generation_draft_context.rs b/crates/router/src/generate/generate_image/image_generation_draft_context.rs index db5b4de5..37155d4c 100644 --- a/crates/router/src/generate/generate_image/image_generation_draft_context.rs +++ b/crates/router/src/generate/generate_image/image_generation_draft_context.rs @@ -3,6 +3,7 @@ use std::fmt::Debug; use sqlite_identifiers::ids::media_file_token::MediaFileToken; +use crate::api::asset_upload_cache::AssetUploadCache; use crate::client::router_client::RouterClient; use crate::client::router_higgsfield_client::RouterHiggsfieldClient; use crate::client::router_seedance2pro_client::RouterSeedance2ProClient; @@ -28,6 +29,10 @@ pub struct ImageGenerationDraftContext<'a> { /// inputs that must be downloaded from ArtCraft and re-uploaded to the /// target provider's CDN. pub media_file_to_artcraft_url_map: Option<&'a HashMap>, + + /// Optional: the account's upload history on the target provider, so + /// media it already holds isn't uploaded again. + pub asset_upload_cache: Option<&'a dyn AssetUploadCache>, } impl<'a> ImageGenerationDraftContext<'a> { @@ -49,6 +54,7 @@ impl Debug for ImageGenerationDraftContext<'_> { f.debug_struct("ImageGenerationDraftContext") .field("client", &self.client.is_some()) .field("media_file_to_artcraft_url_map", &self.media_file_to_artcraft_url_map.map(|m| m.len())) + .field("asset_upload_cache", &self.asset_upload_cache.is_some()) .finish() } } diff --git a/crates/router/src/generate/generate_image/providers/higgsfield/draft.rs b/crates/router/src/generate/generate_image/providers/higgsfield/draft.rs index d5636164..701d22b3 100644 --- a/crates/router/src/generate/generate_image/providers/higgsfield/draft.rs +++ b/crates/router/src/generate/generate_image/providers/higgsfield/draft.rs @@ -5,7 +5,7 @@ use crate::errors::artcraft_router_error::ArtcraftRouterError; use crate::generate::generate_image::image_generation_draft_context::ImageGenerationDraftContext; use crate::generate::generate_image::providers::higgsfield::image_request::HiggsfieldImageRequest; use crate::generate::generate_image::providers::higgsfield::request::HiggsfieldImageRequestState; -use crate::utils::higgsfield_media::{upload_media_list, HiggsfieldMediaKind}; +use crate::utils::higgsfield_media::{upload_media_list, HiggsfieldMediaKind, HiggsfieldUploadContext}; /// A planned Higgsfield image request whose reference images still have to /// be uploaded. `to_request` does the uploads and produces the sendable state. @@ -33,19 +33,21 @@ impl HiggsfieldImageDraftState { draft_context: &ImageGenerationDraftContext<'_>, ) -> Result { let client = draft_context.get_higgsfield_client_ref()?; - let session = &client.session; + let uploads = HiggsfieldUploadContext { + session: &client.session, + maybe_map: draft_context.media_file_to_artcraft_url_map, + maybe_cache: draft_context.asset_upload_cache, + }; let mut reference_images = Vec::new(); if let Some(remaining) = self.unhandled_request_state.take() { - let map = draft_context.media_file_to_artcraft_url_map; // The image generator's references skip the IP check (as the web app // does); only the Seedance video models insist on it. reference_images = upload_media_list( - session, + uploads, remaining.reference_images.map(Into::into), HiggsfieldMediaKind::Image, false, - map, ).await?; info!("Uploaded {} reference image(s) to Higgsfield for {}", reference_images.len(), self.request.model_label()); } diff --git a/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7/request.rs b/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7/request.rs index c04e5c5f..021c700c 100644 --- a/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7/request.rs +++ b/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7/request.rs @@ -129,6 +129,7 @@ mod tests { let ctx = ImageGenerationDraftContext { client: Some(&client), media_file_to_artcraft_url_map: Some(&map), + asset_upload_cache: None, }; let request = ImageGenerationDraftRequest::KinoviMidjourney7(draft) diff --git a/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7_niji/request.rs b/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7_niji/request.rs index 9ea078bc..47c7eef3 100644 --- a/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7_niji/request.rs +++ b/crates/router/src/generate/generate_image/providers/kinovi/midjourney_7_niji/request.rs @@ -127,6 +127,7 @@ mod tests { let ctx = ImageGenerationDraftContext { client: Some(&client), media_file_to_artcraft_url_map: Some(&map), + asset_upload_cache: None, }; let request = ImageGenerationDraftRequest::KinoviMidjourney7Niji(draft) diff --git a/crates/router/src/generate/generate_image/providers/kinovi/midjourney_8/request.rs b/crates/router/src/generate/generate_image/providers/kinovi/midjourney_8/request.rs index 575cfcd9..97107ef0 100644 --- a/crates/router/src/generate/generate_image/providers/kinovi/midjourney_8/request.rs +++ b/crates/router/src/generate/generate_image/providers/kinovi/midjourney_8/request.rs @@ -127,6 +127,7 @@ mod tests { let ctx = ImageGenerationDraftContext { client: Some(&client), media_file_to_artcraft_url_map: Some(&map), + asset_upload_cache: None, }; let request = ImageGenerationDraftRequest::KinoviMidjourney8(draft) diff --git a/crates/router/src/generate/generate_image/providers/kinovi/seedream_5p0_pro/request.rs b/crates/router/src/generate/generate_image/providers/kinovi/seedream_5p0_pro/request.rs index 0c6f7fb8..66f0d609 100644 --- a/crates/router/src/generate/generate_image/providers/kinovi/seedream_5p0_pro/request.rs +++ b/crates/router/src/generate/generate_image/providers/kinovi/seedream_5p0_pro/request.rs @@ -128,6 +128,7 @@ mod tests { let ctx = ImageGenerationDraftContext { client: Some(&client), media_file_to_artcraft_url_map: Some(&map), + asset_upload_cache: None, }; let request = ImageGenerationDraftRequest::KinoviSeedream5p0Pro(draft) diff --git a/crates/router/src/generate/generate_video/providers/higgsfield/draft.rs b/crates/router/src/generate/generate_video/providers/higgsfield/draft.rs index 9d4f59fa..58b81834 100644 --- a/crates/router/src/generate/generate_video/providers/higgsfield/draft.rs +++ b/crates/router/src/generate/generate_video/providers/higgsfield/draft.rs @@ -10,7 +10,7 @@ use crate::generate::generate_video::providers::higgsfield::common::HiggsfieldVi use crate::generate::generate_video::providers::higgsfield::request::HiggsfieldVideoRequestState; use crate::generate::generate_video::providers::higgsfield::video_request::HiggsfieldVideoRequest; use crate::generate::generate_video::video_generation_draft_context::VideoGenerationDraftContext; -use crate::utils::higgsfield_media::{upload_image_ref, upload_media_list, HiggsfieldMediaKind}; +use crate::utils::higgsfield_media::{upload_image_ref, upload_media_list, HiggsfieldMediaKind, HiggsfieldUploadContext}; /// What a `build.rs` planned, before media is uploaded. #[derive(Clone, Debug)] @@ -65,18 +65,21 @@ impl HiggsfieldVideoDraftState { draft_context: &VideoGenerationDraftContext<'_>, ) -> Result { let client = draft_context.get_higgsfield_client_ref()?; - let session = &client.session; - let map = draft_context.media_file_to_artcraft_url_map; + let uploads = HiggsfieldUploadContext { + session: &client.session, + maybe_map: draft_context.media_file_to_artcraft_url_map, + maybe_cache: draft_context.asset_upload_cache, + }; let label = self.plan.model_label(); let references = self.unhandled_request_state.take().unwrap_or_default(); - let start_frame = upload_image_ref(session, references.start_frame, self.ip_check, map).await?; - let end_frame = upload_image_ref(session, references.end_frame, self.ip_check, map).await?; - let images = upload_media_list(session, references.reference_images.map(Into::into), HiggsfieldMediaKind::Image, self.ip_check, map).await?; - let mut videos = upload_media_list(session, references.reference_videos.map(Into::into), HiggsfieldMediaKind::Video, self.ip_check, map).await?; + let start_frame = upload_image_ref(uploads, references.start_frame, self.ip_check).await?; + let end_frame = upload_image_ref(uploads, references.end_frame, self.ip_check).await?; + let images = upload_media_list(uploads, references.reference_images.map(Into::into), HiggsfieldMediaKind::Image, self.ip_check).await?; + let mut videos = upload_media_list(uploads, references.reference_videos.map(Into::into), HiggsfieldMediaKind::Video, self.ip_check).await?; // Audio has no IP check. - let audio = upload_media_list(session, references.reference_audio.map(Into::into), HiggsfieldMediaKind::Audio, false, map).await?; + let audio = upload_media_list(uploads, references.reference_audio.map(Into::into), HiggsfieldMediaKind::Audio, false).await?; let mut request = match self.plan.clone() { HiggsfieldVideoPlan::Request(request) => request, diff --git a/crates/router/src/generate/generate_video/video_generation_draft_context.rs b/crates/router/src/generate/generate_video/video_generation_draft_context.rs index 374e48b8..f0bc54f6 100644 --- a/crates/router/src/generate/generate_video/video_generation_draft_context.rs +++ b/crates/router/src/generate/generate_video/video_generation_draft_context.rs @@ -1,3 +1,4 @@ +use crate::api::asset_upload_cache::AssetUploadCache; use crate::client::router_client::RouterClient; use crate::client::router_higgsfield_client::RouterHiggsfieldClient; use crate::client::router_seedance2pro_client::RouterSeedance2ProClient; @@ -20,6 +21,10 @@ pub struct VideoGenerationDraftContext<'a> { /// Optional context: a map of Character Tokens to their respective Kinovi IDs /// Only necessary if using Kinovi characters pub character_token_to_kinovi_id_map: Option<&'a HashMap>, + + /// Optional: the account's upload history on the target provider, so + /// media it already holds isn't uploaded again. + pub asset_upload_cache: Option<&'a dyn AssetUploadCache>, } impl <'a> VideoGenerationDraftContext<'a> { @@ -52,6 +57,7 @@ impl Debug for VideoGenerationDraftContext<'_> { .field("client", &self.client.is_some()) .field("media_file_to_artcraft_url_map", &self.media_file_to_artcraft_url_map) .field("character_token_to_kinovi_id_map", &self.character_token_to_kinovi_id_map) + .field("asset_upload_cache", &self.asset_upload_cache.is_some()) .finish() } } diff --git a/crates/router/src/utils/higgsfield_media.rs b/crates/router/src/utils/higgsfield_media.rs index 20b60c0b..0c64f41d 100644 --- a/crates/router/src/utils/higgsfield_media.rs +++ b/crates/router/src/utils/higgsfield_media.rs @@ -6,18 +6,26 @@ //! //! Only ArtCraft media tokens touch the ArtCraft cloud (they resolve through //! the token→URL map); local files and bytes go straight to Higgsfield. +//! +//! With an [`AssetUploadCache`] in hand, bytes the account already uploaded +//! are reused by id instead: the cached id is checked against Higgsfield's +//! media status endpoint first, and any doubt means a fresh upload. use std::collections::HashMap; use std::path::Path; +use file_hashing::hash_bytes_blake3; +use higgsfield_client::endpoints::media::get_media_status::MediaStatusFamily; use higgsfield_client::session::higgsfield_session::HiggsfieldSession; use higgsfield_client::session::upload_media::ReferenceMediaFile; use higgsfield_client::session::upload_source_guard::check_upload_source_url; +use higgsfield_client::types::ids::MediaId; use higgsfield_client::types::media_input::MediaInput; use higgsfield_client::types::media_mime_type::MediaMimeType; -use log::info; +use log::{info, warn}; use sqlite_identifiers::ids::media_file_token::MediaFileToken; +use crate::api::asset_upload_cache::{AssetUploadCache, CachedAssetUpload}; use crate::api::audio_list_ref::AudioListRef; use crate::api::image_list_ref::ImageListRef; use crate::api::image_ref::ImageRef; @@ -40,6 +48,24 @@ pub enum HiggsfieldMediaKind { } impl HiggsfieldMediaKind { + /// The status endpoint family for verifying a previous upload. Audio has + /// none, so cached audio ids are trusted as-is. + fn status_family(self) -> Option { + match self { + Self::Image => Some(MediaStatusFamily::Image), + Self::Video => Some(MediaStatusFamily::Video), + Self::Audio => None, + } + } + + fn media_input(self, id: MediaId, url: String) -> MediaInput { + match self { + Self::Image => MediaInput::uploaded(id, url), + Self::Video => MediaInput::uploaded_video(id, url), + Self::Audio => MediaInput::uploaded_audio(id, url), + } + } + fn fallback_mime_type(self) -> MediaMimeType { match self { Self::Image => MediaMimeType::ImageJpeg, @@ -89,14 +115,23 @@ impl HiggsfieldMediaSources { } } +/// Everything an upload needs besides the source itself. +#[derive(Clone, Copy)] +pub struct HiggsfieldUploadContext<'a> { + pub session: &'a HiggsfieldSession, + /// Media file tokens resolve to ArtCraft CDN URLs through this. + pub maybe_map: Option<&'a HashMap>, + /// The account's prior uploads; `None` uploads unconditionally. + pub maybe_cache: Option<&'a dyn AssetUploadCache>, +} + /// Upload every source in `list`, in order. `None` / empty lists upload -/// nothing. Media file tokens are resolved through `maybe_map`. +/// nothing. pub async fn upload_media_list( - session: &HiggsfieldSession, + context: HiggsfieldUploadContext<'_>, list: Option, kind: HiggsfieldMediaKind, ip_check: bool, - maybe_map: Option<&HashMap>, ) -> Result, ArtcraftRouterError> { let sources = match list { None => return Ok(Vec::new()), @@ -104,56 +139,122 @@ pub async fn upload_media_list( }; let mut uploaded = Vec::with_capacity(sources.len()); for source in sources { - uploaded.push(upload_source_to_higgsfield(session, source, kind, ip_check, maybe_map).await?); + uploaded.push(upload_source_to_higgsfield(context, source, kind, ip_check).await?); } Ok(uploaded) } /// Upload a single image reference (a keyframe), if present. pub async fn upload_image_ref( - session: &HiggsfieldSession, + context: HiggsfieldUploadContext<'_>, image_ref: Option, ip_check: bool, - maybe_map: Option<&HashMap>, ) -> Result, ArtcraftRouterError> { match image_ref { None => Ok(None), Some(image_ref) => { - Ok(Some(upload_source_to_higgsfield(session, image_ref.into(), HiggsfieldMediaKind::Image, ip_check, maybe_map).await?)) + Ok(Some(upload_source_to_higgsfield(context, image_ref.into(), HiggsfieldMediaKind::Image, ip_check).await?)) } } } -/// Turn one source into bytes and upload it to Higgsfield as reference -/// media. `ip_check` asks Higgsfield to run (and waits for) its +/// Turn one source into bytes and get a Higgsfield reference for it: the +/// account's previous upload of the same bytes when the cache knows one and +/// Higgsfield confirms it, otherwise a fresh upload (recorded for next +/// time). `ip_check` asks Higgsfield to run (and waits for) its /// intellectual-property check, which the Seedance video models require on -/// images and clips. +/// images and clips; a cached upload that never had it is uploaded again. /// /// The first-party-domain guard applies only to caller-supplied `Url` /// sources: a token that resolves to our own CDN is a deliberate library /// pick, and local paths / bytes never had a URL at all. pub async fn upload_source_to_higgsfield( - session: &HiggsfieldSession, + context: HiggsfieldUploadContext<'_>, source: MediaSourceRef, kind: HiggsfieldMediaKind, ip_check: bool, - maybe_map: Option<&HashMap>, ) -> Result { if let MediaSourceRef::Url(url) = &source { check_upload_source_url(url) .map_err(|err| ArtcraftRouterError::from(ProviderError::Higgsfield(err.into())))?; } - let resolved = resolve_media_source_bytes(source, maybe_map).await?; + let resolved = resolve_media_source_bytes(source, context.maybe_map).await?; + let file_hash = hash_bytes_blake3(&resolved.bytes); + if let Some(cache) = context.maybe_cache { + if let Some(reusable) = reusable_upload(context.session, cache, &file_hash, kind, ip_check).await { + info!("Reusing Higgsfield {:?} reference {} for {} (already uploaded by this account)", kind, reusable.id, resolved.description); + cache.note_reused(&file_hash).await; + return Ok(reusable); + } + } + + let byte_count = resolved.bytes.len() as u64; let file = reference_media_file(resolved, kind, ip_check); info!( "Uploading {:?} reference to Higgsfield ({} bytes, {}, ip_check={})", kind, file.bytes.len(), file.mime_type, ip_check, ); + let uploaded = session_upload(context.session, file).await?; + + if let Some(cache) = context.maybe_cache { + cache.record_upload(&file_hash, byte_count, &CachedAssetUpload { + service_id: uploaded.id.as_str().to_string(), + maybe_service_url: Some(uploaded.url.clone()), + }).await; + } + Ok(uploaded) +} + +async fn session_upload(session: &HiggsfieldSession, file: ReferenceMediaFile) -> Result { session.upload_reference_media(file).await .map_err(|err| ArtcraftRouterError::from(ProviderError::Higgsfield(err))) } +/// The cached upload for `file_hash`, if Higgsfield still honours it. Any +/// problem — no cached URL, a failed status read, a missing or flagged +/// upload, or a required IP check that never ran — forgets the entry and +/// answers `None` so the caller uploads afresh. +async fn reusable_upload( + session: &HiggsfieldSession, + cache: &dyn AssetUploadCache, + file_hash: &str, + kind: HiggsfieldMediaKind, + ip_check: bool, +) -> Option { + let cached = cache.find_upload(file_hash).await?; + let media_id = MediaId::new(cached.service_id.clone()); + + let Some(url) = cached.maybe_service_url.clone() else { + warn!("Cached Higgsfield upload {} has no URL; uploading again", media_id); + cache.forget_upload(file_hash).await; + return None; + }; + + if let Some(family) = kind.status_family() { + match session.get_media_status(family, &media_id).await { + Ok(status) if status.is_ip_detected() => { + warn!("Cached Higgsfield upload {} is flagged as protected content; uploading again", media_id); + cache.forget_upload(file_hash).await; + return None; + } + Ok(status) if ip_check && !status.is_ip_check_finished() => { + info!("Cached Higgsfield upload {} never had its IP check; uploading again with one", media_id); + cache.forget_upload(file_hash).await; + return None; + } + Ok(_) => {} + Err(err) => { + warn!("Could not confirm cached Higgsfield upload {} ({}); uploading again", media_id, err); + cache.forget_upload(file_hash).await; + return None; + } + } + } + + Some(kind.media_input(media_id, url)) +} + /// Describe resolved bytes for upload: the MIME type is sniffed from the /// bytes, then guessed from the name hint (URL, path, or file name), then /// defaulted by `kind`. The resolved source URL (caller-supplied URLs only) @@ -330,6 +431,83 @@ mod tests { } } + mod upload_reuse { + use std::sync::Mutex; + + use async_trait::async_trait; + use higgsfield_client::client::clerk_host::ClerkHost; + use higgsfield_client::client::higgsfield_host::HiggsfieldHost; + use higgsfield_client::types::media_input::MediaInputKind; + + use super::*; + + /// A cache with one canned entry that records what the router did. + struct CannedCache { + entry: Option, + forgotten: Mutex>, + reused: Mutex>, + } + + impl CannedCache { + fn with(entry: CachedAssetUpload) -> Self { + Self { entry: Some(entry), forgotten: Mutex::new(Vec::new()), reused: Mutex::new(Vec::new()) } + } + } + + #[async_trait] + impl AssetUploadCache for CannedCache { + async fn find_upload(&self, _file_hash_blake3: &str) -> Option { + self.entry.clone() + } + async fn record_upload(&self, _file_hash_blake3: &str, _file_size_bytes: u64, _upload: &CachedAssetUpload) {} + async fn note_reused(&self, file_hash_blake3: &str) { + self.reused.lock().unwrap().push(file_hash_blake3.to_string()); + } + async fn forget_upload(&self, file_hash_blake3: &str) { + self.forgotten.lock().unwrap().push(file_hash_blake3.to_string()); + } + } + + /// A session whose every request fails (nothing listens on port 9). + fn unreachable_session() -> HiggsfieldSession { + HiggsfieldSession::from_cookie_header("__client=x") + .with_hosts(HiggsfieldHost::Custom("http://127.0.0.1:9".into()), ClerkHost::Custom("http://127.0.0.1:9".into())) + } + + #[tokio::test] + async fn audio_has_no_status_endpoint_so_the_cached_id_is_trusted() { + let cache = CannedCache::with(CachedAssetUpload { service_id: "media-audio".into(), maybe_service_url: Some("https://cdn/a.mp3".into()) }); + let reused = reusable_upload(&unreachable_session(), &cache, "hash", HiggsfieldMediaKind::Audio, false).await.unwrap(); + assert_eq!(reused.id.as_str(), "media-audio"); + assert_eq!(reused.kind, MediaInputKind::AudioInput); + assert_eq!(reused.url, "https://cdn/a.mp3"); + assert!(cache.forgotten.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn an_entry_without_a_url_is_forgotten() { + let cache = CannedCache::with(CachedAssetUpload { service_id: "media-1".into(), maybe_service_url: None }); + assert!(reusable_upload(&unreachable_session(), &cache, "hash", HiggsfieldMediaKind::Audio, false).await.is_none()); + assert_eq!(*cache.forgotten.lock().unwrap(), vec!["hash".to_string()]); + } + + #[tokio::test] + async fn a_status_check_that_fails_means_a_fresh_upload() { + let cache = CannedCache::with(CachedAssetUpload { service_id: "media-1".into(), maybe_service_url: Some("https://cdn/x.png".into()) }); + assert!(reusable_upload(&unreachable_session(), &cache, "hash", HiggsfieldMediaKind::Image, true).await.is_none()); + assert_eq!(*cache.forgotten.lock().unwrap(), vec!["hash".to_string()]); + assert!(cache.reused.lock().unwrap().is_empty()); + } + + #[test] + fn media_inputs_take_their_kind_from_the_reference_kind() { + let id = MediaId::new("abc"); + assert_eq!(HiggsfieldMediaKind::Image.media_input(id.clone(), "u".into()).kind, MediaInputKind::MediaInput); + assert_eq!(HiggsfieldMediaKind::Video.media_input(id.clone(), "u".into()).kind, MediaInputKind::VideoInput); + assert_eq!(HiggsfieldMediaKind::Audio.media_input(id, "u".into()).kind, MediaInputKind::AudioInput); + } + } + mod token_resolution { use std::collections::HashMap; diff --git a/crates/schema/local_files_database/migrations/20260907000002_create_asset_uploads.sql b/crates/schema/local_files_database/migrations/20260907000002_create_asset_uploads.sql new file mode 100644 index 00000000..52adbbf7 --- /dev/null +++ b/crates/schema/local_files_database/migrations/20260907000002_create_asset_uploads.sql @@ -0,0 +1,37 @@ +-- The upload-dedup ledger: which remote service already holds a copy of a +-- local file, per account. Before uploading, callers look up +-- (service, account, content hash) and reuse the service's id for the +-- asset instead of pushing the bytes again. +-- +-- `service` - which remote: "artcraft", "higgsfield", ... +-- `account_id` - the credential the upload was made with (a credential id +-- for stored accounts; a session fingerprint for ArtCraft +-- web sessions). Different accounts never share uploads. +-- `service_id` - the service's own primary key for the asset (ArtCraft media +-- file token, Higgsfield media id, ...). +-- `service_url` - where the asset can be reached remotely, when the service +-- publishes one (Higgsfield's CDN URL, ArtCraft's CDN URL). +-- +-- Rows are advisory: a reuse is verified against the service first, and any +-- failure falls back to a fresh upload that overwrites the row. + +CREATE TABLE asset_uploads ( + service TEXT NOT NULL, + account_id TEXT NOT NULL, + file_hash_blake3 TEXT NOT NULL, + + service_id TEXT NOT NULL, + service_url TEXT, + + -- Size of the bytes that were uploaded, for sanity checks / display. + file_size_bytes INTEGER, + + uploaded_at TEXT NOT NULL, + -- Last time a caller reused this row instead of uploading. + last_reused_at TEXT, + + PRIMARY KEY (service, account_id, file_hash_blake3) +); + +-- Reverse lookups (which local content is behind a service id). +CREATE INDEX idx_asset_uploads_service_id ON asset_uploads (service, service_id); diff --git a/crates/schema/local_files_database/src/connection.rs b/crates/schema/local_files_database/src/connection.rs index 5a70f7f3..c546d0d7 100644 --- a/crates/schema/local_files_database/src/connection.rs +++ b/crates/schema/local_files_database/src/connection.rs @@ -20,10 +20,10 @@ pub struct LocalFilesDbConnection { impl LocalFilesDbConnection { /// Open (creating if missing) and migrate the database. /// - /// On migration failure the file is deleted and recreated — acceptable - /// while every row is derivable from the filesystem (hashes and thumbnail - /// flags regenerate on demand). Revisit before storing anything - /// irreplaceable here (e.g. the upload-dedup ledger). + /// On migration failure the file is moved aside (`.broken-`) + /// and a fresh one created. Nothing here is irreplaceable — thumbnails + /// regenerate and the upload ledger only costs a re-upload — but the + /// ledger is worth keeping for a hand recovery, so it isn't deleted. pub async fn connect_and_migrate>(database_file: P) -> AnyhowResult { match run_migrations(&database_file).await { Ok(pool) => return Ok(Self { pool }), @@ -32,10 +32,23 @@ impl LocalFilesDbConnection { } } - info!("Deleting and recreating local_files SQLite database at {:?}", database_file.as_ref()); + let database_file = database_file.as_ref(); + let unix_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_millis()) + .unwrap_or(0); + let aside = database_file.with_extension(format!("broken-{unix_ms}")); + info!("Moving local_files SQLite database aside to {:?} and recreating it", aside); - if let Err(err) = std::fs::remove_file(&database_file) { - error!("Error deleting local_files SQLite database file: {:?}", err); + if let Err(err) = std::fs::rename(database_file, &aside) { + error!("Error moving local_files SQLite database file aside: {:?}", err); + } + // WAL/shm sidecars belong to the old file; a stale WAL must not replay + // into the new database. + for sidecar in ["-wal", "-shm"] { + let mut path = database_file.as_os_str().to_owned(); + path.push(sidecar); + let _ = std::fs::remove_file(path); } let pool = run_migrations(database_file).await?; diff --git a/crates/schema/local_files_database/src/queries/asset_upload_record.rs b/crates/schema/local_files_database/src/queries/asset_upload_record.rs new file mode 100644 index 00000000..7182e906 --- /dev/null +++ b/crates/schema/local_files_database/src/queries/asset_upload_record.rs @@ -0,0 +1,15 @@ +use chrono::{DateTime, Utc}; + +/// One row of `asset_uploads`: a remote copy of some local content. +#[derive(Clone, Debug, sqlx::FromRow)] +pub struct AssetUploadRecord { + pub service: String, + pub account_id: String, + pub file_hash_blake3: String, + /// The service's own primary key for the asset. + pub service_id: String, + pub service_url: Option, + pub file_size_bytes: Option, + pub uploaded_at: DateTime, + pub last_reused_at: Option>, +} diff --git a/crates/schema/local_files_database/src/queries/delete_asset_upload.rs b/crates/schema/local_files_database/src/queries/delete_asset_upload.rs new file mode 100644 index 00000000..2dd824e6 --- /dev/null +++ b/crates/schema/local_files_database/src/queries/delete_asset_upload.rs @@ -0,0 +1,21 @@ +use crate::connection::LocalFilesDbConnection; +use crate::error::LocalFilesDbError; + +/// Forget a remote copy (it no longer exists, or failed verification). +/// Returns whether a row was deleted. +pub async fn delete_asset_upload( + db: &LocalFilesDbConnection, + service: &str, + account_id: &str, + file_hash_blake3: &str, +) -> Result { + let result = sqlx::query( + "DELETE FROM asset_uploads WHERE service = ?1 AND account_id = ?2 AND file_hash_blake3 = ?3", + ) + .bind(service) + .bind(account_id) + .bind(file_hash_blake3) + .execute(db.get_pool()) + .await?; + Ok(result.rows_affected() > 0) +} diff --git a/crates/schema/local_files_database/src/queries/get_asset_upload.rs b/crates/schema/local_files_database/src/queries/get_asset_upload.rs new file mode 100644 index 00000000..aba8ed4b --- /dev/null +++ b/crates/schema/local_files_database/src/queries/get_asset_upload.rs @@ -0,0 +1,27 @@ +use crate::connection::LocalFilesDbConnection; +use crate::error::LocalFilesDbError; +use crate::queries::asset_upload_record::AssetUploadRecord; + +pub struct GetAssetUploadArgs<'a> { + pub db: &'a LocalFilesDbConnection, + pub service: &'a str, + pub account_id: &'a str, + pub file_hash_blake3: &'a str, +} + +/// The remote copy of this content on this service+account, if one was +/// recorded. +pub async fn get_asset_upload(args: GetAssetUploadArgs<'_>) -> Result, LocalFilesDbError> { + let record = sqlx::query_as::<_, AssetUploadRecord>( + r#" + SELECT * FROM asset_uploads + WHERE service = ?1 AND account_id = ?2 AND file_hash_blake3 = ?3 + "#, + ) + .bind(args.service) + .bind(args.account_id) + .bind(args.file_hash_blake3) + .fetch_optional(args.db.get_pool()) + .await?; + Ok(record) +} diff --git a/crates/schema/local_files_database/src/queries/mod.rs b/crates/schema/local_files_database/src/queries/mod.rs index ba47600f..00e88009 100644 --- a/crates/schema/local_files_database/src/queries/mod.rs +++ b/crates/schema/local_files_database/src/queries/mod.rs @@ -1,3 +1,4 @@ +// ── local_files: the content-hash index ── pub mod local_file_record; pub mod get_local_file_by_path; pub mod get_local_files_by_hash; @@ -5,3 +6,10 @@ pub mod get_local_files_by_paths; pub mod list_local_files_needing_thumbnails; pub mod set_thumbnail_result; pub mod upsert_local_file; + +// ── asset_uploads: the per-account upload-dedup ledger ── +pub mod asset_upload_record; +pub mod delete_asset_upload; +pub mod get_asset_upload; +pub mod touch_asset_upload_reused; +pub mod upsert_asset_upload; diff --git a/crates/schema/local_files_database/src/queries/touch_asset_upload_reused.rs b/crates/schema/local_files_database/src/queries/touch_asset_upload_reused.rs new file mode 100644 index 00000000..484f603b --- /dev/null +++ b/crates/schema/local_files_database/src/queries/touch_asset_upload_reused.rs @@ -0,0 +1,27 @@ +use chrono::Utc; + +use crate::connection::LocalFilesDbConnection; +use crate::error::LocalFilesDbError; + +/// Note that a caller skipped an upload thanks to this row. Returns whether +/// the row existed. +pub async fn touch_asset_upload_reused( + db: &LocalFilesDbConnection, + service: &str, + account_id: &str, + file_hash_blake3: &str, +) -> Result { + let result = sqlx::query( + r#" + UPDATE asset_uploads SET last_reused_at = ?4 + WHERE service = ?1 AND account_id = ?2 AND file_hash_blake3 = ?3 + "#, + ) + .bind(service) + .bind(account_id) + .bind(file_hash_blake3) + .bind(Utc::now()) + .execute(db.get_pool()) + .await?; + Ok(result.rows_affected() > 0) +} diff --git a/crates/schema/local_files_database/src/queries/upsert_asset_upload.rs b/crates/schema/local_files_database/src/queries/upsert_asset_upload.rs new file mode 100644 index 00000000..fd7ed461 --- /dev/null +++ b/crates/schema/local_files_database/src/queries/upsert_asset_upload.rs @@ -0,0 +1,42 @@ +use chrono::Utc; + +use crate::connection::LocalFilesDbConnection; +use crate::error::LocalFilesDbError; + +pub struct UpsertAssetUploadArgs<'a> { + pub db: &'a LocalFilesDbConnection, + pub service: &'a str, + pub account_id: &'a str, + pub file_hash_blake3: &'a str, + pub service_id: &'a str, + pub maybe_service_url: Option<&'a str>, + pub maybe_file_size_bytes: Option, +} + +/// Record (or replace) the remote copy of some content. A re-upload of the +/// same content overwrites the old service id and resets the reuse time. +pub async fn upsert_asset_upload(args: UpsertAssetUploadArgs<'_>) -> Result<(), LocalFilesDbError> { + sqlx::query( + r#" + INSERT INTO asset_uploads + (service, account_id, file_hash_blake3, service_id, service_url, file_size_bytes, uploaded_at, last_reused_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, NULL) + ON CONFLICT(service, account_id, file_hash_blake3) DO UPDATE SET + service_id = excluded.service_id, + service_url = excluded.service_url, + file_size_bytes = excluded.file_size_bytes, + uploaded_at = excluded.uploaded_at, + last_reused_at = NULL + "#, + ) + .bind(args.service) + .bind(args.account_id) + .bind(args.file_hash_blake3) + .bind(args.service_id) + .bind(args.maybe_service_url) + .bind(args.maybe_file_size_bytes) + .bind(Utc::now()) + .execute(args.db.get_pool()) + .await?; + Ok(()) +} diff --git a/crates/schema/local_files_database/tests/asset_upload_queries.rs b/crates/schema/local_files_database/tests/asset_upload_queries.rs new file mode 100644 index 00000000..c9974ec5 --- /dev/null +++ b/crates/schema/local_files_database/tests/asset_upload_queries.rs @@ -0,0 +1,87 @@ +//! Integration tests for the upload-dedup ledger, against a throwaway +//! migrated database in a tempdir. + +use local_files_database::connection::LocalFilesDbConnection; +use local_files_database::queries::delete_asset_upload::delete_asset_upload; +use local_files_database::queries::get_asset_upload::{get_asset_upload, GetAssetUploadArgs}; +use local_files_database::queries::touch_asset_upload_reused::touch_asset_upload_reused; +use local_files_database::queries::upsert_asset_upload::{upsert_asset_upload, UpsertAssetUploadArgs}; + +const HASH_A: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + +#[tokio::test] +async fn record_lookup_and_reuse() { + let dir = tempfile::tempdir().unwrap(); + let db = test_db(&dir).await; + + assert!(lookup(&db, "higgsfield", "credential_1", HASH_A).await.is_none()); + + record(&db, "higgsfield", "credential_1", HASH_A, "media-1", Some("https://cdn.higgsfield.ai/x.png")).await; + let row = lookup(&db, "higgsfield", "credential_1", HASH_A).await.unwrap(); + assert_eq!(row.service_id, "media-1"); + assert_eq!(row.service_url.as_deref(), Some("https://cdn.higgsfield.ai/x.png")); + assert_eq!(row.file_size_bytes, Some(1234)); + assert!(row.last_reused_at.is_none()); + + assert!(touch_asset_upload_reused(&db, "higgsfield", "credential_1", HASH_A).await.unwrap()); + assert!(lookup(&db, "higgsfield", "credential_1", HASH_A).await.unwrap().last_reused_at.is_some()); + assert!(!touch_asset_upload_reused(&db, "higgsfield", "nobody", HASH_A).await.unwrap()); +} + +#[tokio::test] +async fn accounts_and_services_are_isolated() { + let dir = tempfile::tempdir().unwrap(); + let db = test_db(&dir).await; + + record(&db, "higgsfield", "credential_1", HASH_A, "media-1", None).await; + record(&db, "higgsfield", "credential_2", HASH_A, "media-2", None).await; + record(&db, "artcraft", "credential_1", HASH_A, "m_token", None).await; + + assert_eq!(lookup(&db, "higgsfield", "credential_1", HASH_A).await.unwrap().service_id, "media-1"); + assert_eq!(lookup(&db, "higgsfield", "credential_2", HASH_A).await.unwrap().service_id, "media-2"); + assert_eq!(lookup(&db, "artcraft", "credential_1", HASH_A).await.unwrap().service_id, "m_token"); + assert!(lookup(&db, "artcraft", "credential_2", HASH_A).await.is_none()); +} + +#[tokio::test] +async fn reupload_replaces_and_delete_forgets() { + let dir = tempfile::tempdir().unwrap(); + let db = test_db(&dir).await; + + record(&db, "higgsfield", "credential_1", HASH_A, "media-old", Some("https://old")).await; + touch_asset_upload_reused(&db, "higgsfield", "credential_1", HASH_A).await.unwrap(); + + record(&db, "higgsfield", "credential_1", HASH_A, "media-new", None).await; + let row = lookup(&db, "higgsfield", "credential_1", HASH_A).await.unwrap(); + assert_eq!(row.service_id, "media-new"); + assert!(row.service_url.is_none()); + assert!(row.last_reused_at.is_none(), "a fresh upload resets the reuse marker"); + + assert!(delete_asset_upload(&db, "higgsfield", "credential_1", HASH_A).await.unwrap()); + assert!(!delete_asset_upload(&db, "higgsfield", "credential_1", HASH_A).await.unwrap()); + assert!(lookup(&db, "higgsfield", "credential_1", HASH_A).await.is_none()); +} + +// ── Helpers ── + +async fn test_db(dir: &tempfile::TempDir) -> LocalFilesDbConnection { + LocalFilesDbConnection::connect_and_migrate(dir.path().join("local_files.sqlite")) + .await + .expect("connect and migrate") +} + +async fn record(db: &LocalFilesDbConnection, service: &str, account_id: &str, hash: &str, service_id: &str, url: Option<&str>) { + upsert_asset_upload(UpsertAssetUploadArgs { + db, + service, + account_id, + file_hash_blake3: hash, + service_id, + maybe_service_url: url, + maybe_file_size_bytes: Some(1234), + }).await.expect("upsert"); +} + +async fn lookup(db: &LocalFilesDbConnection, service: &str, account_id: &str, hash: &str) -> Option { + get_asset_upload(GetAssetUploadArgs { db, service, account_id, file_hash_blake3: hash }).await.expect("get") +} diff --git a/frontend/apps/artcraft/app/src/state/usePromptboxStatePersistence.ts b/frontend/apps/artcraft/app/src/state/usePromptboxStatePersistence.ts index b3e675c6..fa8c0d61 100644 --- a/frontend/apps/artcraft/app/src/state/usePromptboxStatePersistence.ts +++ b/frontend/apps/artcraft/app/src/state/usePromptboxStatePersistence.ts @@ -187,9 +187,8 @@ const pickOptions = (store: Record, spec: FieldSpec): Recordaccount memory. -const snapshot = (): UpdatePromptboxStateRequest[] => { +// The complete state to persist, as the backend stores it. +const snapshot = (): UpdatePromptboxStateRequest => { const selection = useClassyModelSelectorStore.getState(); const accountId = selection.selectedAccountId ?? undefined; const optionsFor: Partial>> = { @@ -198,41 +197,39 @@ const snapshot = (): UpdatePromptboxStateRequest[] => { audio: pickOptions(asRecord(usePromptAudioStore.getState()), AUDIO_OPTION_FIELDS), }; - const patches: UpdatePromptboxStateRequest[] = SELECTOR_PAGES.map(({ modality, page }) => ({ - modality, + const modalityState = (modality: PromptboxModality, selectedModel: string | undefined): ModalityPromptboxState => ({ selected_account_id: accountId, - selected_model: selection.selectedModels[page]?.id, + selected_model: selectedModel, options: optionsFor[modality] ?? {}, - })); - patches.push({ - modality: "audio", - selected_account_id: accountId, - selected_model: usePromptAudioStore.getState().selectedModelId ?? undefined, - options: optionsFor.audio ?? {}, }); - patches.push({ last_account_by_model: selection.lastAccountByModel }); - return patches; + const selectorModel = (page: ModelPage) => selection.selectedModels[page]?.id; + + return { + image: modalityState("image", selectorModel(ModelPage.TextToImage)), + video: modalityState("video", selectorModel(ModelPage.ImageToVideo)), + splat: modalityState("splat", selectorModel(ModelPage.ImageTo3DWorld)), + mesh: modalityState("mesh", selectorModel(ModelPage.ImageTo3DObject)), + audio: modalityState("audio", usePromptAudioStore.getState().selectedModelId ?? undefined), + last_account_by_model: selection.lastAccountByModel, + }; }; -// Subscribe to every store that feeds the snapshot; write only the patches -// that changed since the last write, debounced. +// Subscribe to every store that feeds the snapshot; write the whole snapshot +// (debounced) whenever it differs from the last one written. const startPersisting = (): (() => void) => { - const lastWritten = new Map(); - for (const patch of snapshot()) lastWritten.set(patchKey(patch), JSON.stringify(patch)); + let lastWritten = JSON.stringify(snapshot()); let timer: number | undefined; const flush = () => { timer = undefined; - for (const patch of snapshot()) { - const key = patchKey(patch); - const json = JSON.stringify(patch); - if (lastWritten.get(key) === json) continue; - lastWritten.set(key, json); - UpdatePromptboxState(patch).catch((err) => { - console.warn("[promptbox] could not persist state:", err); - lastWritten.delete(key); // Try again on the next change. - }); - } + const state = snapshot(); + const json = JSON.stringify(state); + if (json === lastWritten) return; + lastWritten = json; + UpdatePromptboxState(state).catch((err) => { + console.warn("[promptbox] could not persist state:", err); + lastWritten = ""; // Try again on the next change. + }); }; const schedule = () => { if (timer !== undefined) window.clearTimeout(timer); @@ -250,6 +247,3 @@ const startPersisting = (): (() => void) => { for (const unsubscribe of unsubscribers) unsubscribe(); }; }; - -const patchKey = (patch: UpdatePromptboxStateRequest): string => - patch.modality ?? "last_account_by_model"; diff --git a/frontend/libs/tauri-api/src/lib/promptbox/PromptboxState.ts b/frontend/libs/tauri-api/src/lib/promptbox/PromptboxState.ts index 2b654737..df12a022 100644 --- a/frontend/libs/tauri-api/src/lib/promptbox/PromptboxState.ts +++ b/frontend/libs/tauri-api/src/lib/promptbox/PromptboxState.ts @@ -33,19 +33,13 @@ export interface GetPromptboxStateResult { export const GetPromptboxState = async (): Promise => (await invoke("get_promptbox_state_command")) as GetPromptboxStateResult; -// A patch: only present fields change. `options` replaces the modality's -// options wholesale; `last_account_by_model` entries merge in. -export interface UpdatePromptboxStateRequest { - modality?: PromptboxModality; - selected_account_id?: string; - selected_model?: string; - options?: Record; - last_account_by_model?: Record; -} +// The full state to persist. The backend replaces what it has wholesale and +// owns `version`, so the frontend never sends one. +export type UpdatePromptboxStateRequest = Omit; export interface UpdatePromptboxStateResult { state: PromptboxState; } -export const UpdatePromptboxState = async (request: UpdatePromptboxStateRequest): Promise => - (await invoke("update_promptbox_state_command", { request })) as UpdatePromptboxStateResult; +export const UpdatePromptboxState = async (state: UpdatePromptboxStateRequest): Promise => + (await invoke("update_promptbox_state_command", { state })) as UpdatePromptboxStateResult;