finish up endpoints, plumb events through tauri

This commit is contained in:
Brandon Thomas committed 2025-08-21 01:28:29 -04:00
1 parent e8e4c1298e
commit 72d148ed7a
13 files changed
+213 -45

No files matched your search

@@ -0,0 +1,19 @@
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
use crate::error::storyteller_error::StorytellerError;
use crate::utils::api_host::ApiHost;
use crate::utils::basic_json_get_request::basic_json_get_request;
use artcraft_api_defs::media_file::list_batch_generated_media_files::{ListBatchGeneratedReduxMediaFilesSuccessResponse, LIST_BATCH_GENERATED_REDUX_MEDIA_FILES_URL_PATH};
use tokens::tokens::batch_generations::BatchGenerationToken;
pub async fn list_batch_generated_redux_media_files(
api_host: &ApiHost,
maybe_creds: Option<&StorytellerCredentialSet>,
batch_token: &BatchGenerationToken,
) -> Result<ListBatchGeneratedReduxMediaFilesSuccessResponse, StorytellerError> {
let url = format!("{}/{}", LIST_BATCH_GENERATED_REDUX_MEDIA_FILES_URL_PATH, batch_token.to_string());
Ok(basic_json_get_request(
api_host,
&url,
maybe_creds,
).await?)
}
@@ -1,4 +1,5 @@
pub mod get_media_file;
pub mod list_batch_generated_redux_media_files;
pub mod upload_image_media_file_from_bytes;
pub mod upload_image_media_file_from_file;
pub mod upload_new_engine_asset_from_file;
@@ -0,0 +1,27 @@
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
use enums::tauri::ux::tauri_event_name::TauriEventName;
use serde_derive::Serialize;
use tokens::tokens::media_files::MediaFileToken;
use url::Url;
#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "snake_case")]
pub struct ImageEditCompleteEvent {
// pub source_media_token: MediaFileToken,
pub edited_images: Vec<EditedImage>,
pub maybe_frontend_subscriber_id: Option<String>,
pub maybe_frontend_subscriber_payload: Option<String>,
}
#[derive(Clone, Debug, Serialize)]
#[serde(rename_all = "snake_case")]
pub struct EditedImage {
pub media_token: MediaFileToken,
pub cdn_url: Url,
pub maybe_thumbnail_template: Option<String>,
}
impl BasicSendableEvent for ImageEditCompleteEvent {
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::CanvasBgRemovedEvent;
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success;
}
@@ -1,3 +1,4 @@
pub mod canvas_background_removal_complete_event;
pub mod image_edit_complete_event;
pub mod refresh_account_state_event;
pub mod show_provider_login_modal_event;
@@ -1,16 +1,25 @@
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent;
use crate::core::events::functional_events::image_edit_complete_event::{EditedImage, ImageEditCompleteEvent};
use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
use anyhow::anyhow;
use artcraft_api_defs::jobs::list_session_jobs::{ListSessionJobsItem, ListSessionResultDetailsResponse};
use enums::tauri::tasks::task_type::TaskType;
use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
use errors::AnyhowResult;
use log::error;
use log::{error, warn};
use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
use storyteller_client::media_files::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files;
use tauri::AppHandle;
use tokens::tokens::batch_generations::BatchGenerationToken;
use tokens::tokens::media_files::MediaFileToken;
pub async fn maybe_handle_inpainting_complete_event(
app: &tauri::AppHandle,
app: &AppHandle,
app_env_configs: &AppEnvConfigs,
maybe_creds: Option<&StorytellerCredentialSet>,
task: &Task,
job: &ListSessionJobsItem,
) -> AnyhowResult<()> {
@@ -25,7 +34,7 @@ pub async fn maybe_handle_inpainting_complete_event(
_ => return Ok(()),
}
let result = match job.maybe_result {
let job_result = match job.maybe_result {
Some(ref res) => res,
None => {
error!("Job result is None for task: {:?}", task);
@@ -33,18 +42,84 @@ pub async fn maybe_handle_inpainting_complete_event(
},
};
// TODO: Lookup all media files in the batch.
let event = CanvasBackgroundRemovalCompleteEvent {
media_token: MediaFileToken::new_from_str(&result.entity_token),
image_cdn_url: result.media_links.cdn_url.clone(),
maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(),
maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(),
let event = match job_result.maybe_batch_token.as_ref() {
Some(batch_token) => {
handle_batch(
app,
app_env_configs,
maybe_creds,
task,
job,
job_result,
batch_token
).await?
}
None => {
handle_single(
app,
task,
job,
job_result,
).await?
}
};
if let Err(err) = event.send(&app) {
error!("Failed to send CanvasBackgroundRemovalCompleteEvent: {:?}", err); // Fail open
error!("Failed to send ImageEditCompleteEvent: {:?}", err); // Fail open
}
Ok(())
}
pub async fn handle_single(
app: &AppHandle,
task: &Task,
job: &ListSessionJobsItem,
job_result: &ListSessionResultDetailsResponse,
) -> AnyhowResult<ImageEditCompleteEvent> {
Ok(ImageEditCompleteEvent {
edited_images: vec![EditedImage {
media_token: MediaFileToken::new_from_str(&job_result.entity_token),
cdn_url: job_result.media_links.cdn_url.clone(),
maybe_thumbnail_template: job_result.media_links.maybe_thumbnail_template.clone(),
}],
maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(),
maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(),
})
}
pub async fn handle_batch(
app: &AppHandle,
app_env_configs: &AppEnvConfigs,
maybe_creds: Option<&StorytellerCredentialSet>,
task: &Task,
job: &ListSessionJobsItem,
job_result: &ListSessionResultDetailsResponse,
batch_token: &BatchGenerationToken,
) -> AnyhowResult<ImageEditCompleteEvent> {
let result = list_batch_generated_redux_media_files(
&app_env_configs.storyteller_host,
maybe_creds,
batch_token,
).await?;
if result.media_files.is_empty() {
return Err(anyhow!("No media files found for batch token: {}", batch_token));
}
let media_files = result.media_files
.into_iter()
.map(|file| EditedImage {
media_token: file.token,
cdn_url: file.media_links.cdn_url,
maybe_thumbnail_template: file.media_links.maybe_thumbnail_template,
})
.collect();
Ok(ImageEditCompleteEvent {
edited_images: media_files,
maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(),
maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(),
})
}
@@ -33,6 +33,9 @@ use storyteller_client::jobs::list_session_jobs::{list_session_jobs, States};
use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
use tauri::AppHandle;
use tempdir::TempDir;
use artcraft_api_defs::jobs::list_session_jobs::ListSessionJobsItem;
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
use crate::services::storyteller::threads::events::maybe_handle_inpainting_complete_event::maybe_handle_inpainting_complete_event;
pub async fn storyteller_task_polling_thread(
app_handle: AppHandle,
@@ -140,16 +143,14 @@ async fn polling_loop(
continue; // If anything breaks with queries, don't spam events.
}
let result = maybe_send_background_removal_complete_event(
send_additional_events(
&app_handle,
&task,
&app_env_configs,
creds.as_ref(),
&job,
&task,
).await;
if let Err(err) = result {
error!("Failed to send background removal complete event: {:?}", err);
}
let service = to_generation_service_provider(task.provider);
let action = to_generation_action(task.task_type);
@@ -166,3 +167,33 @@ async fn polling_loop(
}
}
async fn send_additional_events(
app_handle: &AppHandle,
app_env_configs: &AppEnvConfigs,
creds: Option<&StorytellerCredentialSet>,
job: &ListSessionJobsItem,
task: &Task
) {
let result = maybe_handle_inpainting_complete_event(
app_handle,
app_env_configs,
creds,
task,
job,
).await;
if let Err(err) = result {
error!("Failed to send image inpainting complete event: {:?}", err);
}
let result = maybe_send_background_removal_complete_event(
app_handle,
task,
job,
).await;
if let Err(err) = result {
error!("Failed to send background removal complete event: {:?}", err);
}
}
@@ -119,31 +119,31 @@ pub struct MediaFileRaw {
pub updated_at: DateTime<Utc>,
}
pub async fn list_batch_generated_media_files(
pub async fn list_batch_generated_redux_media_files(
batch_token: &BatchGenerationToken,
can_see_deleted: bool,
mysql_pool: &MySqlPool
) -> AnyhowResult<MediaFileBatch> {
list_batch_generated_media_files_with_transactor(
list_batch_generated_redux_media_files_with_transactor(
batch_token,
can_see_deleted,
Transactor::for_pool(mysql_pool)
).await
}
pub async fn list_batch_generated_media_files_with_connection(
pub async fn list_batch_generated_redux_media_files_with_connection(
batch_token: &BatchGenerationToken,
can_see_deleted: bool,
mysql_connection: &mut PoolConnection<MySql>,
) -> AnyhowResult<MediaFileBatch> {
list_batch_generated_media_files_with_transactor(
list_batch_generated_redux_media_files_with_transactor(
batch_token,
can_see_deleted,
Transactor::for_connection(mysql_connection)
).await
}
pub async fn list_batch_generated_media_files_with_transactor(
pub async fn list_batch_generated_redux_media_files_with_transactor(
batch_token: &BatchGenerationToken,
can_see_deleted: bool,
transactor: Transactor<'_, '_>,
@@ -1,4 +1,4 @@
pub mod list_batch_generated_media_files;
pub mod list_batch_generated_redux_media_files;
pub mod list_featured_media_files;
pub mod list_media_files;
pub mod list_media_files_by_batch_token;
@@ -11,20 +11,22 @@ use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::prompts::PromptToken;
use utoipa::ToSchema;
pub const LIST_BATCH_GENERATED_REDUX_MEDIA_FILES_URL_PATH: &str = "/v1/media_files/batch_gen_redux";
/// For the URL PathInfo
#[derive(Serialize, Deserialize, ToSchema)]
pub struct ListBatchGeneratedMediaFilesPathInfo {
pub struct ListBatchGeneratedReduxMediaFilesPathInfo {
pub token: BatchGenerationToken,
}
#[derive(Serialize, Deserialize, ToSchema)]
pub struct ListBatchGeneratedMediaFilesSuccessResponse {
pub struct ListBatchGeneratedReduxMediaFilesSuccessResponse {
pub success: bool,
pub media_files: Vec<BatchGeneratedMediaFileInfo>,
pub media_files: Vec<BatchGeneratedReduxMediaFileInfo>,
}
#[derive(Serialize, Deserialize, ToSchema)]
pub struct BatchGeneratedMediaFileInfo {
pub struct BatchGeneratedReduxMediaFileInfo {
/// Primary key identifier
pub token: MediaFileToken,
@@ -43,6 +43,11 @@ pub enum TauriEventName {
#[serde(rename = "canvas_bg_removed_event")]
CanvasBgRemovedEvent,
/// Special event:
/// Image edit is complete
#[serde(rename = "image_edit_complete_event")]
ImageEditCompleteEvent,
/// Special event:
/// Refresh account states
#[serde(rename = "refresh_account_state_event")]
@@ -73,6 +78,7 @@ impl TauriEventName {
Self::GenerationCompleteEvent => "generation-complete-event",
Self::GenerationFailedEvent => "generation-failed-event",
Self::CanvasBgRemovedEvent => "canvas_bg_removed_event",
Self::ImageEditCompleteEvent => "image_edit_complete_event",
Self::RefreshAccountStateEvent => "refresh_account_state_event",
Self::ShowProviderLoginModalEvent => "show_provider_login_modal_event",
Self::FlashUserInputErrorEvent => "flash_user_input_error_event",
@@ -86,6 +92,7 @@ impl TauriEventName {
"generation-complete-event" => Ok(Self::GenerationCompleteEvent),
"generation-failed-event" => Ok(Self::GenerationFailedEvent),
"canvas_bg_removed_event" => Ok(Self::CanvasBgRemovedEvent),
"image_edit_complete_event" => Ok(Self::ImageEditCompleteEvent),
"refresh_account_state_event" => Ok(Self::RefreshAccountStateEvent),
"show_provider_login_modal_event" => Ok(Self::ShowProviderLoginModalEvent),
"flash_user_input_error_event" => Ok(Self::FlashUserInputErrorEvent),
@@ -102,6 +109,7 @@ impl TauriEventName {
Self::GenerationCompleteEvent,
Self::GenerationFailedEvent,
Self::CanvasBgRemovedEvent,
Self::ImageEditCompleteEvent,
Self::RefreshAccountStateEvent,
Self::ShowProviderLoginModalEvent,
Self::FlashUserInputErrorEvent,
@@ -124,6 +132,7 @@ mod tests {
assert_serialization(TauriEventName::GenerationCompleteEvent, "generation-complete-event");
assert_serialization(TauriEventName::GenerationFailedEvent, "generation-failed-event");
assert_serialization(TauriEventName::CanvasBgRemovedEvent, "canvas_bg_removed_event");
assert_serialization(TauriEventName::ImageEditCompleteEvent, "image_edit_complete_event");
assert_serialization(TauriEventName::RefreshAccountStateEvent, "refresh_account_state_event");
assert_serialization(TauriEventName::ShowProviderLoginModalEvent, "show_provider_login_modal_event");
assert_serialization(TauriEventName::FlashUserInputErrorEvent, "flash_user_input_error_event");
@@ -136,6 +145,7 @@ mod tests {
assert_eq!(TauriEventName::GenerationCompleteEvent.to_str(), "generation-complete-event");
assert_eq!(TauriEventName::GenerationFailedEvent.to_str(), "generation-failed-event");
assert_eq!(TauriEventName::CanvasBgRemovedEvent.to_str(), "canvas_bg_removed_event");
assert_eq!(TauriEventName::ImageEditCompleteEvent.to_str(), "image_edit_complete_event");
assert_eq!(TauriEventName::RefreshAccountStateEvent.to_str(), "refresh_account_state_event");
assert_eq!(TauriEventName::ShowProviderLoginModalEvent.to_str(), "show_provider_login_modal_event");
assert_eq!(TauriEventName::FlashUserInputErrorEvent.to_str(), "flash_user_input_error_event");
@@ -148,6 +158,7 @@ mod tests {
assert_eq!(TauriEventName::from_str("generation-complete-event").unwrap(), TauriEventName::GenerationCompleteEvent);
assert_eq!(TauriEventName::from_str("generation-failed-event").unwrap(), TauriEventName::GenerationFailedEvent);
assert_eq!(TauriEventName::from_str("canvas_bg_removed_event").unwrap(), TauriEventName::CanvasBgRemovedEvent);
assert_eq!(TauriEventName::from_str("image_edit_complete_event").unwrap(), TauriEventName::ImageEditCompleteEvent);
assert_eq!(TauriEventName::from_str("refresh_account_state_event").unwrap(), TauriEventName::RefreshAccountStateEvent);
assert_eq!(TauriEventName::from_str("show_provider_login_modal_event").unwrap(), TauriEventName::ShowProviderLoginModalEvent);
assert_eq!(TauriEventName::from_str("flash_user_input_error_event").unwrap(), TauriEventName::FlashUserInputErrorEvent);
@@ -156,12 +167,13 @@ mod tests {
#[test]
fn all_variants() {
let mut variants = TauriEventName::all_variants();
assert_eq!(variants.len(), 8);
assert_eq!(variants.len(), 9);
assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueSuccessEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueFailureEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationCompleteEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationFailedEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::CanvasBgRemovedEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::ImageEditCompleteEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::RefreshAccountStateEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::ShowProviderLoginModalEvent));
assert_eq!(variants.pop_first(), Some(TauriEventName::FlashUserInputErrorEvent));
@@ -18,7 +18,7 @@ use actix_web::{web, HttpRequest, HttpResponse};
use actix_web_lab::extract::Query;
use artcraft_api_defs::common::responses::media_links::MediaLinks;
use artcraft_api_defs::common::responses::simple_entity_stats::SimpleEntityStats;
use artcraft_api_defs::media_file::list_batch_generated_media_files::{BatchGeneratedMediaFileInfo, ListBatchGeneratedMediaFilesPathInfo, ListBatchGeneratedMediaFilesSuccessResponse};
use artcraft_api_defs::media_file::list_batch_generated_media_files::{BatchGeneratedReduxMediaFileInfo, ListBatchGeneratedReduxMediaFilesPathInfo, ListBatchGeneratedReduxMediaFilesSuccessResponse};
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use chrono::{DateTime, Utc};
use enums::by_table::media_files::media_file_animation_type::MediaFileAnimationType;
@@ -33,30 +33,30 @@ use enums::no_table::style_transfer::style_transfer_name::StyleTransferName;
use enums_public::by_table::model_weights::public_weights_types::PublicWeightsType;
use log::warn;
use mysql_queries::queries::media_files::get::batch_get_media_files::batch_get_media_files;
use mysql_queries::queries::media_files::list::list_batch_generated_media_files::list_batch_generated_media_files_with_connection;
use mysql_queries::queries::media_files::list::list_batch_generated_redux_media_files::list_batch_generated_redux_media_files_with_connection;
use tokens::tokens::batch_generations::BatchGenerationToken;
use tokens::tokens::media_files::MediaFileToken;
use tokens::tokens::model_weights::ModelWeightToken;
use tokens::tokens::prompts::PromptToken;
use utoipa::{IntoParams, ToSchema};
/// List media files generated as part of the same batch
/// List media files generated as part of the same batch (redux endpoint)
#[utoipa::path(
get,
tag = "Media Files",
path = "/v1/media_files/batch_generated/{token}",
path = "/v1/media_files/batch_gen_redux/{token}",
responses(
(status = 200, description = "Success", body = ListBatchGeneratedMediaFilesSuccessResponse),
(status = 200, description = "Success", body = ListBatchGeneratedReduxMediaFilesSuccessResponse),
),
params(
("request" = ListBatchGeneratedMediaFilesPathInfo, description = "Payload for Request"),
("request" = ListBatchGeneratedReduxMediaFilesPathInfo, description = "Payload for Request"),
)
)]
pub async fn list_batch_generated_media_files_handler(
pub async fn list_batch_generated_redux_media_files_handler(
http_request: HttpRequest,
path: Path<ListBatchGeneratedMediaFilesPathInfo>,
path: Path<ListBatchGeneratedReduxMediaFilesPathInfo>,
server_state: web::Data<Arc<ServerState>>
) -> Result<Json<ListBatchGeneratedMediaFilesSuccessResponse>, CommonWebError> {
) -> Result<Json<ListBatchGeneratedReduxMediaFilesSuccessResponse>, CommonWebError> {
let mut mysql_connection = server_state.mysql_pool
.acquire()
.await?;
@@ -84,7 +84,7 @@ pub async fn list_batch_generated_media_files_handler(
// NB(bt,2024-03-24): I'm sorry, this is gross. We're not respecting sorting, input ordering,
// de-duplication, (if we swap types), etc. Gotta go fast.
let result = list_batch_generated_media_files_with_connection(
let result = list_batch_generated_redux_media_files_with_connection(
&path.token,
show_deleted_results,
&mut mysql_connection,
@@ -107,7 +107,7 @@ pub async fn list_batch_generated_media_files_handler(
result.maybe_public_bucket_prefix.as_deref(),
result.maybe_public_bucket_extension.as_deref());
BatchGeneratedMediaFileInfo {
BatchGeneratedReduxMediaFileInfo {
token: result.token.clone(),
media_class: result.media_class,
media_type: result.media_type,
@@ -137,7 +137,7 @@ pub async fn list_batch_generated_media_files_handler(
})
.collect();
Ok(Json(ListBatchGeneratedMediaFilesSuccessResponse {
Ok(Json(ListBatchGeneratedReduxMediaFilesSuccessResponse {
success: true,
media_files,
}))
@@ -1,4 +1,4 @@
pub mod list_batch_generated_media_files_handler;
pub mod list_batch_generated_redux_media_files_handler;
pub mod list_featured_media_files_handler;
pub mod list_media_files_by_batch_token_handler;
pub mod list_media_files_for_user_handler;
@@ -12,7 +12,7 @@ use crate::http_server::endpoints::media_files::edit::set_media_file_cover_image
use crate::http_server::endpoints::media_files::edit::update_media_file_handler::update_media_file_handler;
use crate::http_server::endpoints::media_files::get::batch_get_media_files_handler::batch_get_media_files_handler;
use crate::http_server::endpoints::media_files::get::get_media_file_handler::get_media_file_handler;
use crate::http_server::endpoints::media_files::list::list_batch_generated_media_files_handler::list_batch_generated_media_files_handler;
use crate::http_server::endpoints::media_files::list::list_batch_generated_redux_media_files_handler::list_batch_generated_redux_media_files_handler;
use crate::http_server::endpoints::media_files::list::list_featured_media_files_handler::list_featured_media_files_handler;
use crate::http_server::endpoints::media_files::list::list_media_files_by_batch_token_handler::list_media_files_by_batch_token_handler;
use crate::http_server::endpoints::media_files::list::list_media_files_for_user_handler::list_media_files_for_user_handler;
@@ -77,8 +77,8 @@ pub fn add_media_file_routes<T, B> (app: App<T>) -> App<T>
.route(web::get().to(batch_get_media_files_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(web::resource("/batch_generated/{token}")
.route(web::get().to(list_batch_generated_media_files_handler))
.service(web::resource("/batch_gen_redux/{token}")
.route(web::get().to(list_batch_generated_redux_media_files_handler))
.route(web::head().to(|| HttpResponse::Ok()))
)
.service(web::resource("/list")