more endpoints

crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4_edit_image_webhook.rs

@@ -1,6 +1,7 @@
 use crate::creds::fal_api_key::FalApiKey;
 use crate::error::classify_fal_error::classify_fal_error;
 use crate::error::fal_error_plus::FalErrorPlus;
+use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents};
 use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4_edit_image::{bytedance_seedream_v4_edit_image, BytedanceSeedreamV4EditImageInput};
 use fal::webhook::WebhookResponse;
 use reqwest::IntoUrl;
@@ -53,6 +54,23 @@ pub enum EnqueueBytedanceSeedreamV4EditImageSize {
   Auto4k,
 }
 
+
+impl <U: IntoUrl> FalRequestCostCalculator for EnqueueBytedanceSeedreamV4EditImageArgs<'_, U> {
+  fn calculate_cost_in_cents(&self) -> UsdCents {
+    // Your request will cost $0.03 per image.
+    let unit_cost = 3;
+    let cost = match self.num_images {
+      None => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4EditImageNumImages::One) => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4EditImageNumImages::Two) => unit_cost * 2,
+      Some(EnqueueBytedanceSeedreamV4EditImageNumImages::Three) => unit_cost * 3,
+      Some(EnqueueBytedanceSeedreamV4EditImageNumImages::Four) => unit_cost * 4,
+    };
+    cost as UsdCents
+  }
+}
+
+
 pub async fn enqueue_bytedance_seedream_v4_edit_image_webhook<R: IntoUrl>(
   args: EnqueueBytedanceSeedreamV4EditImageArgs<'_, R>
 ) -> Result<WebhookResponse, FalErrorPlus> {

crates/api_clients/fal_client/src/requests/webhook/image/edit/enqueue_bytedance_seedream_v4p5_edit_image_webhook.rs

@@ -1,6 +1,7 @@
 use crate::creds::fal_api_key::FalApiKey;
 use crate::error::classify_fal_error::classify_fal_error;
 use crate::error::fal_error_plus::FalErrorPlus;
+use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents};
 use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4p5_edit_image::{bytedance_seedream_v4p5_edit_image, BytedanceSeedreamV4p5EditImageInput};
 use fal::webhook::WebhookResponse;
 use reqwest::IntoUrl;
@@ -52,6 +53,23 @@ pub enum EnqueueBytedanceSeedreamV4p5EditImageSize {
   Auto4k,
 }
 
+
+impl <U: IntoUrl> FalRequestCostCalculator for EnqueueBytedanceSeedreamV4p5EditImageArgs<'_, U> {
+  fn calculate_cost_in_cents(&self) -> UsdCents {
+    // Your request will cost $0.04 per image.
+    let unit_cost = 4;
+    let cost = match self.num_images {
+      None => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::One) => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::Two) => unit_cost * 2,
+      Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::Three) => unit_cost * 3,
+      Some(EnqueueBytedanceSeedreamV4p5EditImageNumImages::Four) => unit_cost * 4,
+    };
+    cost as UsdCents
+  }
+}
+
+
 pub async fn enqueue_bytedance_seedream_v4p5_edit_image_webhook<R: IntoUrl>(
   args: EnqueueBytedanceSeedreamV4p5EditImageArgs<'_, R>
 ) -> Result<WebhookResponse, FalErrorPlus> {

crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4_text_to_image_webhook.rs

@@ -1,6 +1,7 @@
 use crate::creds::fal_api_key::FalApiKey;
 use crate::error::classify_fal_error::classify_fal_error;
 use crate::error::fal_error_plus::FalErrorPlus;
+use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents};
 use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4_text_to_image::BytedanceSeedreamV4TextToImageInput;
 use fal::prelude::fal_ai::bytedance::seedream::bytedance_seedream_v4_text_to_image::bytedance_seedream_v4_text_to_image;
 use fal::webhook::WebhookResponse;
@@ -53,6 +54,24 @@ pub enum EnqueueBytedanceSeedreamV4TextToImageSize {
   Auto4k,
 }
 
+
+impl <U: IntoUrl> FalRequestCostCalculator for EnqueueBytedanceSeedreamV4TextToImageArgs<'_, U> {
+  fn calculate_cost_in_cents(&self) -> UsdCents {
+    // Copied from edit image version
+    // Your request will cost $0.03 per image.
+    let unit_cost = 3;
+    let cost = match self.num_images {
+      None => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::One) => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::Two) => unit_cost * 2,
+      Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::Three) => unit_cost * 3,
+      Some(EnqueueBytedanceSeedreamV4TextToImageNumImages::Four) => unit_cost * 4,
+    };
+    cost as UsdCents
+  }
+}
+
+
 pub async fn enqueue_bytedance_seedream_v4_text_to_image_webhook<R: IntoUrl>(
   args: EnqueueBytedanceSeedreamV4TextToImageArgs<'_, R>
 ) -> Result<WebhookResponse, FalErrorPlus> {

crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_bytedance_seedream_v4p5_text_to_image_webhook.rs

@@ -4,6 +4,8 @@ use crate::error::fal_error_plus::FalErrorPlus;
 use fal::endpoints::fal_ai::bytedance::seedream::bytedance_seedream_v4p5_text_to_image::{bytedance_seedream_v4p5_text_to_image, BytedanceSeedreamV4p5TextToImageInput};
 use fal::webhook::WebhookResponse;
 use reqwest::IntoUrl;
+use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents};
+use crate::requests::webhook::image::edit::enqueue_bytedance_seedream_v4p5_edit_image_webhook::{EnqueueBytedanceSeedreamV4p5EditImageArgs, EnqueueBytedanceSeedreamV4p5EditImageNumImages};
 
 pub struct EnqueueBytedanceSeedreamV4p5TextToImageArgs<'a, R: IntoUrl> {
   // Request required
@@ -51,6 +53,24 @@ pub enum EnqueueBytedanceSeedreamV4p5TextToImageSize {
   Auto4k,
 }
 
+
+impl <U: IntoUrl> FalRequestCostCalculator for EnqueueBytedanceSeedreamV4p5TextToImageArgs<'_, U> {
+  fn calculate_cost_in_cents(&self) -> UsdCents {
+    // NB: Copied from edit image costs
+    // Your request will cost $0.04 per image.
+    let unit_cost = 4;
+    let cost = match self.num_images {
+      None => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::One) => unit_cost,
+      Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::Two) => unit_cost * 2,
+      Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::Three) => unit_cost * 3,
+      Some(EnqueueBytedanceSeedreamV4p5TextToImageNumImages::Four) => unit_cost * 4,
+    };
+    cost as UsdCents
+  }
+}
+
+
 pub async fn enqueue_bytedance_seedream_v4p5_text_to_image_webhook<R: IntoUrl>(
   args: EnqueueBytedanceSeedreamV4p5TextToImageArgs<'_, R>
 ) -> Result<WebhookResponse, FalErrorPlus> {

crates/api_clients/fal_client/src/requests/webhook/image/text/enqueue_gemini_25_flash_text_to_image_webhook.rs

@@ -5,6 +5,8 @@ use fal::prelude::fal_ai::gemini_25_flash_image::text_to_image::gemini_25_flash_
 use fal::prelude::fal_ai::gemini_25_flash_image::text_to_image::Gemini25FlashTextToImageInput;
 use fal::webhook::WebhookResponse;
 use reqwest::IntoUrl;
+use crate::requests::traits::fal_request_cost_calculator_trait::{FalRequestCostCalculator, UsdCents};
+use crate::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::{Gemini25FlashEditArgs, Gemini25FlashEditNumImages};
 
 pub struct Gemini25FlashTextToImageArgs<'a, R: IntoUrl> {
   // Request required
@@ -46,6 +48,21 @@ pub enum Gemini25FlashTextToImageAspectRatio {
   NineBySixteen, // NB: No NineByTwentyOne ?
 }
 
+
+impl <U: IntoUrl> FalRequestCostCalculator for Gemini25FlashTextToImageArgs<'_, U> {
+  fn calculate_cost_in_cents(&self) -> UsdCents {
+    // NB: This was copied from the image-to-image case.
+    // Your request will cost $0.039 per image. For $1.00, you can run this model 25 times.
+    match self.num_images {
+      Gemini25FlashTextToImageNumImages::One => 4,
+      Gemini25FlashTextToImageNumImages::Two => 8,
+      Gemini25FlashTextToImageNumImages::Three => 12,
+      Gemini25FlashTextToImageNumImages::Four => 16,
+    }
+  }
+}
+
+
 pub async fn enqueue_gemini_25_flash_text_to_image_webhook<R: IntoUrl>(
   args: Gemini25FlashTextToImageArgs<'_, R>
 ) -> Result<WebhookResponse, FalErrorPlus> {

crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4_multi_function_image_gen_handler.rs

@@ -2,6 +2,7 @@ use std::collections::HashSet;
 use std::iter::FromIterator;
 use std::sync::Arc;
 
+use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
 use crate::http_server::common_responses::common_web_error::CommonWebError;
 use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder;
 use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
@@ -19,6 +20,7 @@ use enums::common::generation_provider::GenerationProvider;
 use enums::common::model_type::ModelType;
 use enums::common::visibility::Visibility;
 use fal_client::creds::open_ai_api_key::OpenAiApiKey;
+use fal_client::requests::traits::fal_request_cost_calculator_trait::FalRequestCostCalculator;
 use fal_client::requests::webhook::image::edit::enqueue_bytedance_seedream_v4_edit_image_webhook::{enqueue_bytedance_seedream_v4_edit_image_webhook, EnqueueBytedanceSeedreamV4EditImageArgs, EnqueueBytedanceSeedreamV4EditImageNumImages, EnqueueBytedanceSeedreamV4EditImageSize};
 use fal_client::requests::webhook::image::text::enqueue_bytedance_seedream_v4_text_to_image_webhook::{enqueue_bytedance_seedream_v4_text_to_image_webhook, EnqueueBytedanceSeedreamV4TextToImageArgs, EnqueueBytedanceSeedreamV4TextToImageNumImages, EnqueueBytedanceSeedreamV4TextToImageSize};
 use http_server_common::request::get_request_ip::get_request_ip;
@@ -26,6 +28,7 @@ use log::{error, info, warn};
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue;
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory;
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs;
+use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{insert_generic_inference_job_for_fal_queue_with_apriori_job_token, InsertGenericInferenceForFalWithAprioriJobTokenArgs};
 use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
 use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection};
 use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem};
@@ -33,6 +36,7 @@ use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPrompt
 use server_environment::ServerEnvironment;
 use sqlx::pool::PoolConnection;
 use sqlx::{Acquire, MySql};
+use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::media_files::MediaFileToken;
 use utoipa::ToSchema;
 
@@ -73,15 +77,12 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler(
       .avt_cookie_manager
       .get_avt_token_from_request(&http_request);
 
-  // TODO: Limit usage for new accounts. Billing, free credits metering, etc.
-
-  //let user_session = match maybe_user_session {
-  //  Some(session) => session,
-  //  None => {
-  //    warn!("not logged in");
-  //    return Err(CommonWebError::NotAuthorized);
-  //  }
-  //};
+  let user_token = match maybe_user_session.as_ref() {
+    Some(session) => &session.user_token,
+    None => {
+      return Err(CommonWebError::NotAuthorized);
+    }
+  };
 
   if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
     return Err(CommonWebError::BadInputWithSimpleMessage(reason));
@@ -109,6 +110,8 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler(
 
   info!("Fal webhook URL: {}", server_state.fal.webhook_url);
 
+  let apriori_job_token = InferenceJobToken::generate();
+
   let fal_result;
 
   if let Some(input_image_urls) = image_urls.as_deref() {
@@ -149,6 +152,17 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler(
       api_key: &server_state.fal.api_key,
     };
 
+    let cost = args.calculate_cost_in_cents();
+
+    info!("Charging wallet: {}", cost);
+
+    attempt_wallet_deduction_else_common_web_error(
+      user_token,
+      Some(apriori_job_token.as_str()),
+      cost,
+      &mut mysql_connection,
+    ).await?;
+
     fal_result = enqueue_bytedance_seedream_v4_edit_image_webhook(args)
         .await
         .map_err(|err| {
@@ -193,6 +207,17 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler(
       api_key: &server_state.fal.api_key,
     };
 
+    let cost = args.calculate_cost_in_cents();
+
+    info!("Charging wallet: {}", cost);
+
+    attempt_wallet_deduction_else_common_web_error(
+      user_token,
+      Some(apriori_job_token.as_str()),
+      cost,
+      &mut mysql_connection,
+    ).await?;
+
     fal_result = enqueue_bytedance_seedream_v4_text_to_image_webhook(args)
         .await
         .map_err(|err| {
@@ -223,9 +248,7 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler(
   let prompt_result = insert_prompt(InsertPromptArgs {
     maybe_apriori_prompt_token: None,
     prompt_type: PromptType::ArtcraftApp,
-    maybe_creator_user_token: maybe_user_session
-        .as_ref()
-        .map(|s| &s.user_token),
+    maybe_creator_user_token: Some(&user_token),
     maybe_model_type: Some(ModelType::Seedream4),
     maybe_generation_provider: Some(GenerationProvider::Artcraft),
     maybe_positive_prompt: request.prompt.as_deref(),
@@ -264,13 +287,14 @@ pub async fn bytedance_seedream_v4_multi_function_image_gen_handler(
     }
   }
 
-  let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs {
+  let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(InsertGenericInferenceForFalWithAprioriJobTokenArgs {
+    apriori_job_token: &apriori_job_token,
     uuid_idempotency_token: &request.uuid_idempotency_token,
     maybe_external_third_party_id: &external_job_id,
     fal_category: FalCategory::ImageGeneration,
     maybe_inference_args: None,
     maybe_prompt_token: prompt_token.as_ref(),
-    maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token),
+    maybe_creator_user_token: Some(&user_token),
     maybe_avt_token: maybe_avt_token.as_ref(),
     creator_ip_address: &ip_address,
     creator_set_visibility: Visibility::Public,

crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/bytedance_seedream_v4p5_multi_function_image_gen_handler.rs

@@ -2,6 +2,7 @@ use std::collections::HashSet;
 use std::iter::FromIterator;
 use std::sync::Arc;
 
+use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
 use crate::http_server::common_responses::common_web_error::CommonWebError;
 use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder;
 use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
@@ -19,6 +20,7 @@ use enums::common::generation_provider::GenerationProvider;
 use enums::common::model_type::ModelType;
 use enums::common::visibility::Visibility;
 use fal_client::creds::open_ai_api_key::OpenAiApiKey;
+use fal_client::requests::traits::fal_request_cost_calculator_trait::FalRequestCostCalculator;
 use fal_client::requests::webhook::image::edit::enqueue_bytedance_seedream_v4p5_edit_image_webhook::{enqueue_bytedance_seedream_v4p5_edit_image_webhook, EnqueueBytedanceSeedreamV4p5EditImageArgs, EnqueueBytedanceSeedreamV4p5EditImageNumImages, EnqueueBytedanceSeedreamV4p5EditImageSize};
 use fal_client::requests::webhook::image::text::enqueue_bytedance_seedream_v4p5_text_to_image_webhook::{enqueue_bytedance_seedream_v4p5_text_to_image_webhook, EnqueueBytedanceSeedreamV4p5TextToImageArgs, EnqueueBytedanceSeedreamV4p5TextToImageNumImages, EnqueueBytedanceSeedreamV4p5TextToImageSize};
 use http_server_common::request::get_request_ip::get_request_ip;
@@ -26,6 +28,7 @@ use log::{error, info, warn};
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue;
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory;
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs;
+use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{insert_generic_inference_job_for_fal_queue_with_apriori_job_token, InsertGenericInferenceForFalWithAprioriJobTokenArgs};
 use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
 use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection};
 use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem};
@@ -33,6 +36,7 @@ use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPrompt
 use server_environment::ServerEnvironment;
 use sqlx::pool::PoolConnection;
 use sqlx::{Acquire, MySql};
+use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::media_files::MediaFileToken;
 use utoipa::ToSchema;
 
@@ -73,15 +77,12 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler(
       .avt_cookie_manager
       .get_avt_token_from_request(&http_request);
 
-  // TODO: Limit usage for new accounts. Billing, free credits metering, etc.
-
-  //let user_session = match maybe_user_session {
-  //  Some(session) => session,
-  //  None => {
-  //    warn!("not logged in");
-  //    return Err(CommonWebError::NotAuthorized);
-  //  }
-  //};
+  let user_token = match maybe_user_session.as_ref() {
+    Some(session) => &session.user_token,
+    None => {
+      return Err(CommonWebError::NotAuthorized);
+    }
+  };
 
   if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
     return Err(CommonWebError::BadInputWithSimpleMessage(reason));
@@ -109,6 +110,8 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler(
 
   info!("Fal webhook URL: {}", server_state.fal.webhook_url);
 
+  let apriori_job_token = InferenceJobToken::generate();
+
   let fal_result;
 
   if let Some(input_image_urls) = image_urls.as_deref() {
@@ -148,6 +151,17 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler(
       api_key: &server_state.fal.api_key,
     };
 
+    let cost = args.calculate_cost_in_cents();
+
+    info!("Charging wallet: {}", cost);
+
+    attempt_wallet_deduction_else_common_web_error(
+      user_token,
+      Some(apriori_job_token.as_str()),
+      cost,
+      &mut mysql_connection,
+    ).await?;
+
     fal_result = enqueue_bytedance_seedream_v4p5_edit_image_webhook(args)
         .await
         .map_err(|err| {
@@ -191,6 +205,17 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler(
       api_key: &server_state.fal.api_key,
     };
 
+    let cost = args.calculate_cost_in_cents();
+
+    info!("Charging wallet: {}", cost);
+
+    attempt_wallet_deduction_else_common_web_error(
+      user_token,
+      Some(apriori_job_token.as_str()),
+      cost,
+      &mut mysql_connection,
+    ).await?;
+
     fal_result = enqueue_bytedance_seedream_v4p5_text_to_image_webhook(args)
         .await
         .map_err(|err| {
@@ -221,9 +246,7 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler(
   let prompt_result = insert_prompt(InsertPromptArgs {
     maybe_apriori_prompt_token: None,
     prompt_type: PromptType::ArtcraftApp,
-    maybe_creator_user_token: maybe_user_session
-        .as_ref()
-        .map(|s| &s.user_token),
+    maybe_creator_user_token: Some(&user_token),
     maybe_model_type: Some(ModelType::Seedream4p5),
     maybe_generation_provider: Some(GenerationProvider::Artcraft),
     maybe_positive_prompt: request.prompt.as_deref(),
@@ -262,13 +285,14 @@ pub async fn bytedance_seedream_v4p5_multi_function_image_gen_handler(
     }
   }
 
-  let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs {
+  let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(InsertGenericInferenceForFalWithAprioriJobTokenArgs {
+    apriori_job_token: &apriori_job_token,
     uuid_idempotency_token: &request.uuid_idempotency_token,
     maybe_external_third_party_id: &external_job_id,
     fal_category: FalCategory::ImageGeneration,
     maybe_inference_args: None,
     maybe_prompt_token: prompt_token.as_ref(),
-    maybe_creator_user_token: maybe_user_session.as_ref().map(|s| &s.user_token),
+    maybe_creator_user_token: Some(&user_token),
     maybe_avt_token: maybe_avt_token.as_ref(),
     creator_ip_address: &ip_address,
     creator_set_visibility: Visibility::Public,

crates/service/web/storyteller_web/src/http_server/endpoints/generate/image/multi_function/nano_banana_multi_function_image_gen_handler.rs

@@ -2,6 +2,7 @@ use std::collections::HashSet;
 use std::iter::FromIterator;
 use std::sync::Arc;
 
+use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error;
 use crate::http_server::common_responses::common_web_error::CommonWebError;
 use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder;
 use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
@@ -22,6 +23,7 @@ use enums::common::payments_namespace::PaymentsNamespace;
 use enums::common::stripe_subscription_status::StripeSubscriptionStatus;
 use enums::common::visibility::Visibility;
 use fal_client::creds::open_ai_api_key::OpenAiApiKey;
+use fal_client::requests::traits::fal_request_cost_calculator_trait::FalRequestCostCalculator;
 use fal_client::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::{enqueue_gemini_25_flash_edit_webhook, Gemini25FlashEditArgs, Gemini25FlashEditAspectRatio, Gemini25FlashEditNumImages};
 use fal_client::requests::webhook::image::edit::enqueue_nano_banana_pro_edit_image_webhook::{enqueue_nano_banana_pro_image_edit_webhook, EnqueueNanoBananaProEditImageArgs, EnqueueNanoBananaProEditImageAspectRatio, EnqueueNanoBananaProEditImageNumImages, EnqueueNanoBananaProEditImageResolution};
 use fal_client::requests::webhook::image::text::enqueue_gemini_25_flash_text_to_image_webhook::{enqueue_gemini_25_flash_text_to_image_webhook, Gemini25FlashTextToImageArgs, Gemini25FlashTextToImageAspectRatio, Gemini25FlashTextToImageNumImages};
@@ -31,6 +33,7 @@ use log::{error, info, warn};
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::insert_generic_inference_job_for_fal_queue;
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory;
 use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs;
+use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{insert_generic_inference_job_for_fal_queue_with_apriori_job_token, InsertGenericInferenceForFalWithAprioriJobTokenArgs};
 use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
 use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_connection};
 use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem};
@@ -39,6 +42,7 @@ use mysql_queries::queries::users::user_subscriptions::find_subscription_for_own
 use server_environment::ServerEnvironment;
 use sqlx::pool::PoolConnection;
 use sqlx::{Acquire, MySql};
+use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::media_files::MediaFileToken;
 use utoipa::ToSchema;
 
@@ -83,39 +87,31 @@ pub async fn nano_banana_multi_function_image_gen_handler(
       .avt_cookie_manager
       .get_avt_token_from_request(&http_request);
 
-  let maybe_user_token = maybe_user_session
-      .as_ref()
-      .map(|s| &s.user_token);
-
-  // TODO: Limit usage for new accounts. Billing, free credits metering, etc.
-
-  //let user_session = match maybe_user_session {
-  //  Some(session) => session,
-  //  None => {
-  //    warn!("not logged in");
-  //    return Err(CommonWebError::NotAuthorized);
-  //  }
-  //};
+  let user_token = match maybe_user_session.as_ref() {
+    Some(session) => &session.user_token,
+    None => {
+      return Err(CommonWebError::NotAuthorized);
+    }
+  };
 
   let mut downgrade_for_free_user = true;
 
-  if let Some(user_token) = maybe_user_token.as_ref() {
-    let result = find_subscription_for_owner_user_using_connection(
-      user_token,
-      PaymentsNamespace::Artcraft,
-      &mut mysql_connection,
-    ).await;
+  let result = find_subscription_for_owner_user_using_connection(
+    user_token,
+    PaymentsNamespace::Artcraft,
+    &mut mysql_connection,
+  ).await;
 
-    if let Ok(Some(subscription)) = result {
-      info!("User {:?} has subscription: {:?} (stripe customer: {:?}, status: {:?})", 
-        user_token,
-        subscription.token, 
-        subscription.stripe_customer_id,
-        subscription.stripe_subscription_status);
-      // NB: Failing open means subscribers might get fewer results, but they're free right now.
-      if subscription.stripe_subscription_status == StripeSubscriptionStatus::Active {
-        downgrade_for_free_user = false;
-      }
+  if let Ok(Some(subscription)) = result {
+    info!("User {:?} has subscription: {:?} (stripe customer: {:?}, status: {:?})",
+      user_token,
+      subscription.token,
+      subscription.stripe_customer_id,
+      subscription.stripe_subscription_status);
+
+    // NB: Failing open means subscribers might get fewer results, but they're free right now.
+    if subscription.stripe_subscription_status == StripeSubscriptionStatus::Active {
+      downgrade_for_free_user = false;
     }
   }
 
@@ -143,6 +139,8 @@ pub async fn nano_banana_multi_function_image_gen_handler(
 
   info!("Fal webhook URL: {}", server_state.fal.webhook_url);
 
+  let apriori_job_token = InferenceJobToken::generate();
+
   let fal_result;
 
   if let Some(input_image_urls) = image_urls.as_deref() {
@@ -184,6 +182,17 @@ pub async fn nano_banana_multi_function_image_gen_handler(
       api_key: &server_state.fal.api_key,
     };
 
+    let cost = args.calculate_cost_in_cents();
+
+    info!("Charging wallet: {}", cost);
+
+    attempt_wallet_deduction_else_common_web_error(
+      user_token,
+      Some(apriori_job_token.as_str()),
+      cost,
+      &mut mysql_connection,
+    ).await?;
+
     fal_result = enqueue_gemini_25_flash_edit_webhook(args)
         .await
         .map_err(|err| {
@@ -229,6 +238,17 @@ pub async fn nano_banana_multi_function_image_gen_handler(
       api_key: &server_state.fal.api_key,
     };
 
+    let cost = args.calculate_cost_in_cents();
+
+    info!("Charging wallet: {}", cost);
+
+    attempt_wallet_deduction_else_common_web_error(
+      user_token,
+      Some(apriori_job_token.as_str()),
+      cost,
+      &mut mysql_connection,
+    ).await?;
+
     fal_result = enqueue_gemini_25_flash_text_to_image_webhook(args)
         .await
         .map_err(|err| {
@@ -259,7 +279,7 @@ pub async fn nano_banana_multi_function_image_gen_handler(
   let prompt_result = insert_prompt(InsertPromptArgs {
     maybe_apriori_prompt_token: None,
     prompt_type: PromptType::ArtcraftApp,
-    maybe_creator_user_token: maybe_user_token,
+    maybe_creator_user_token: Some(&user_token),
     maybe_model_type: Some(ModelType::NanoBanana),
     maybe_generation_provider: Some(GenerationProvider::Artcraft),
     maybe_positive_prompt: request.prompt.as_deref(),
@@ -298,13 +318,14 @@ pub async fn nano_banana_multi_function_image_gen_handler(
     }
   }
 
-  let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs {
+  let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(InsertGenericInferenceForFalWithAprioriJobTokenArgs {
+    apriori_job_token: &apriori_job_token,
     uuid_idempotency_token: &request.uuid_idempotency_token,
     maybe_external_third_party_id: &external_job_id,
     fal_category: FalCategory::ImageGeneration,
     maybe_inference_args: None,
     maybe_prompt_token: prompt_token.as_ref(),
-    maybe_creator_user_token: maybe_user_token,
+    maybe_creator_user_token: Some(&user_token),
     maybe_avt_token: maybe_avt_token.as_ref(),
     creator_ip_address: &ip_address,
     creator_set_visibility: Visibility::Public,