update handler

crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/video/omni_gen_video_generate_handler.rs

@@ -38,7 +38,7 @@ use tokens::tokens::media_files::MediaFileToken;
 
 use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
 use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test;
-use crate::http_server::endpoints::omni_gen::generate::video::helpers::bill_wallet::bill_wallet;
+use crate::http_server::endpoints::omni_gen::generate::video::helpers::bill_wallet::{bill_wallet, BillWalletResult};
 use crate::http_server::endpoints::omni_gen::generate::video::pipeline_v1::distill_helper::hydrate_to_router_request::hydrate_to_router_request;
 use crate::http_server::endpoints::omni_gen::generate::video::pipeline_v1::distill_video_request::distill_video_request;
 use crate::http_server::endpoints::omni_gen::generate::video::pipeline_v1::execute::execute_generation::execute_generation;
@@ -48,6 +48,12 @@ use crate::http_server::validations::validate_idempotency_token_format::validate
 use crate::state::server_state::ServerState;
 use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map;
 
+/// Both pipelines produce this, then the shared suffix handles DB writes.
+struct PipelineResult {
+  billing: BillWalletResult,
+  response: GenerateVideoResponse,
+}
+
 /// Generate a video using the omni-gen unified endpoint.
 #[utoipa::path(
   post,
@@ -159,118 +165,57 @@ pub async fn omni_gen_video_generate_handler(
     }
   };
 
-  // ==================== PIPELINE DECISION ==================== //
+  // ==================== PIPELINE DISPATCH ==================== //
 
   let use_v2 = matches!(
     request.model,
     Some(CommonVideoModel::Seedance2p0) | Some(CommonVideoModel::Seedance2p0Fast)
   );
 
-  if use_v2 {
-    execute_v2_pipeline(
-      &http_request,
-      &request,
-      &server_state,
-      &mut mysql_connection,
-      user_token,
-      maybe_avt_token.as_ref(),
-      &idempotency_token,
-      maybe_prompt_model_type,
-      media_file_hydration_map,
-    ).await
+  let pipeline_result = if use_v2 {
+    run_pipeline_v2(&request, &server_state, &mut mysql_connection, user_token, &media_file_hydration_map).await?
   } else {
-    execute_v1_pipeline(
-      &http_request,
-      &request,
-      &server_state,
-      &mut mysql_connection,
-      user_token,
-      maybe_avt_token.as_ref(),
-      &idempotency_token,
-      maybe_prompt_model_type,
-      media_file_hydration_map,
-    ).await
-  }
-}
-
-// ==================== PIPELINE V1 (existing) ==================== //
-
-#[allow(clippy::too_many_arguments)]
-async fn execute_v1_pipeline(
-  http_request: &HttpRequest,
-  request: &OmniGenVideoCostAndGenerateRequest,
-  server_state: &ServerState,
-  mysql_connection: &mut sqlx::pool::PoolConnection<sqlx::MySql>,
-  user_token: &tokens::tokens::users::UserToken,
-  maybe_avt_token: Option<&tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken>,
-  idempotency_token: &str,
-  maybe_prompt_model_type: Option<CommonModelType>,
-  media_file_hydration_map: Option<HashMap<MediaFileToken, Url>>,
-) -> Result<Json<OmniGenVideoGenerateResponse>, AdvancedCommonWebError> {
-  let kinovi_character_ids = resolve_kinovi_character_ids(
-    request.reference_character_tokens.as_deref(),
-    mysql_connection,
-  ).await?;
-
-  let execution_provider = match request.model {
-    Some(CommonVideoModel::Seedance2p0) => Provider::Seedance2Pro,
-    Some(CommonVideoModel::Seedance2p0Fast) => Provider::Seedance2Pro,
-    _ => Provider::Fal,
+    run_pipeline_v1(&request, &server_state, &mut mysql_connection, user_token, &media_file_hydration_map).await?
   };
 
-  let distilled = distill_video_request(request, media_file_hydration_map.as_ref(), execution_provider)?;
+  // ==================== SHARED: DB WRITES ==================== //
 
-  info!("v1 distilled plan: {:?}", distilled.plan);
+  let ip_address = get_request_ip(&http_request);
 
-  let cost = distilled.cost.cost_in_credits.unwrap_or(0);
-  let billing = bill_wallet(user_token, cost, mysql_connection).await?;
-
-  let gen_result = execute_generation(
-    &distilled,
-    request,
-    server_state,
-    media_file_hydration_map.as_ref(),
-    kinovi_character_ids,
-    billing.maybe_wallet_ledger_entry_token.as_ref(),
-    mysql_connection,
-  ).await?;
-
-  // -- DB writes --
-  let ip_address = get_request_ip(http_request);
   let mut transaction = mysql_connection.begin().await.map_err(|err| {
     error!("Error starting MySQL transaction: {:?}", err);
     AdvancedCommonWebError::from_error(err)
   })?;
 
-  let prompt_token = insert_prompt_and_context(
-    request, &mut transaction, user_token, maybe_prompt_model_type,
-    Some(gen_result.generation_mode), &ip_address,
+  insert_prompt_and_context(
+    &request, &mut transaction, user_token, maybe_prompt_model_type, &ip_address,
   ).await;
 
-  let job_token = if gen_result.is_seedance2pro {
-    insert_seedance2pro_jobs(
-      &gen_result.external_job_id,
-      gen_result.maybe_seedance_order_ids.as_deref(),
-      &billing.apriori_job_token,
-      idempotency_token,
-      prompt_token.as_ref(),
-      billing.maybe_wallet_ledger_entry_token.as_ref(),
-      user_token,
-      maybe_avt_token,
-      &ip_address,
-      &mut transaction,
-    ).await?
-  } else {
-    insert_fal_job(
-      &gen_result.external_job_id,
-      &billing.apriori_job_token,
-      idempotency_token,
-      prompt_token.as_ref(),
-      user_token,
-      maybe_avt_token,
-      &ip_address,
-      &mut transaction,
-    ).await?
+  let job_token = match &pipeline_result.response {
+    GenerateVideoResponse::Seedance2Pro(payload) => {
+      insert_seedance2pro_jobs(
+        &payload.order_id, None,
+        &pipeline_result.billing.apriori_job_token, &idempotency_token,
+        pipeline_result.billing.maybe_wallet_ledger_entry_token.as_ref(),
+        user_token, maybe_avt_token.as_ref(), &ip_address,
+        &mut transaction,
+      ).await?
+    }
+    GenerateVideoResponse::Fal(payload) => {
+      let external_id = payload.request_id.clone().unwrap_or_default();
+      insert_fal_job(
+        &external_id, &pipeline_result.billing.apriori_job_token, &idempotency_token,
+        user_token, maybe_avt_token.as_ref(), &ip_address,
+        &mut transaction,
+      ).await?
+    }
+    GenerateVideoResponse::Artcraft(payload) => {
+      payload.inference_job_token.clone()
+    }
+    other => {
+      error!("Unexpected generation response variant: {:?}", other);
+      return Err(AdvancedCommonWebError::server_error_with_message("Unexpected generation response"));
+    }
   };
 
   transaction.commit().await.map_err(|err| {
@@ -284,21 +229,67 @@ async fn execute_v1_pipeline(
   }))
 }
 
-// ==================== PIPELINE V2 (new) ==================== //
+// ==================== PIPELINE V1 ==================== //
 
-#[allow(clippy::too_many_arguments)]
-async fn execute_v2_pipeline(
-  http_request: &HttpRequest,
+async fn run_pipeline_v1(
   request: &OmniGenVideoCostAndGenerateRequest,
   server_state: &ServerState,
   mysql_connection: &mut sqlx::pool::PoolConnection<sqlx::MySql>,
   user_token: &tokens::tokens::users::UserToken,
-  maybe_avt_token: Option<&tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken>,
-  idempotency_token: &str,
-  maybe_prompt_model_type: Option<CommonModelType>,
-  media_file_hydration_map: Option<HashMap<MediaFileToken, Url>>,
-) -> Result<Json<OmniGenVideoGenerateResponse>, AdvancedCommonWebError> {
-  // 1. Build execution request (provider = Seedance2Pro for Kinovi execution)
+  media_file_hydration_map: &Option<HashMap<MediaFileToken, Url>>,
+) -> Result<PipelineResult, AdvancedCommonWebError> {
+  let kinovi_character_ids = resolve_kinovi_character_ids(
+    request.reference_character_tokens.as_deref(),
+    mysql_connection,
+  ).await?;
+
+  let execution_provider = match request.model {
+    Some(CommonVideoModel::Seedance2p0) => Provider::Seedance2Pro,
+    Some(CommonVideoModel::Seedance2p0Fast) => Provider::Seedance2Pro,
+    _ => Provider::Fal,
+  };
+
+  let distilled = distill_video_request(request, media_file_hydration_map.as_ref(), execution_provider)?;
+  info!("v1 distilled plan: {:?}", distilled.plan);
+
+  let cost = distilled.cost.cost_in_credits.unwrap_or(0);
+  let billing = bill_wallet(user_token, cost, mysql_connection).await?;
+
+  let gen_result = execute_generation(
+    &distilled, request, server_state, media_file_hydration_map.as_ref(),
+    kinovi_character_ids, billing.maybe_wallet_ledger_entry_token.as_ref(), mysql_connection,
+  ).await?;
+
+  // Map v1 GenerationResult → GenerateVideoResponse for the shared suffix
+  let response = if gen_result.is_seedance2pro {
+    GenerateVideoResponse::Seedance2Pro(
+      artcraft_router::generate::generate_video::generate_video_response::Seedance2proVideoResponsePayload {
+        order_id: gen_result.external_job_id,
+        task_id: String::new(),
+      }
+    )
+  } else {
+    GenerateVideoResponse::Fal(
+      artcraft_router::generate::generate_video::generate_video_response::FalVideoResponsePayload {
+        request_id: Some(gen_result.external_job_id),
+        gateway_request_id: None,
+      }
+    )
+  };
+
+  Ok(PipelineResult { billing, response })
+}
+
+// ==================== PIPELINE V2 ==================== //
+
+async fn run_pipeline_v2(
+  request: &OmniGenVideoCostAndGenerateRequest,
+  server_state: &ServerState,
+  mysql_connection: &mut sqlx::pool::PoolConnection<sqlx::MySql>,
+  user_token: &tokens::tokens::users::UserToken,
+  media_file_hydration_map: &Option<HashMap<MediaFileToken, Url>>,
+) -> Result<PipelineResult, AdvancedCommonWebError> {
+  // 1. Build execution request (provider = Seedance2Pro for Kinovi)
   let mut exec_builder = hydrate_to_router_request(request)?;
   exec_builder.provider = Provider::Seedance2Pro;
 
@@ -311,18 +302,19 @@ async fn execute_v2_pipeline(
   let mut cost_builder = hydrate_to_router_request(request)?;
   cost_builder.provider = Provider::Artcraft;
 
-  let cost_estimate = cost_builder.build2()
+  let cost = cost_builder.build2()
     .map_err(|e| {
-      warn!("Failed to build2 cost estimate for v2 pipeline: {}", e);
+      warn!("Failed to build2 cost estimate for v2: {}", e);
       AdvancedCommonWebError::from_error(e)
     })?
     .estimate_cost()
     .map_err(|e| {
-      warn!("Failed to estimate cost for v2 pipeline: {}", e);
+      warn!("Failed to estimate cost for v2: {}", e);
       AdvancedCommonWebError::from_error(e)
-    })?;
+    })?
+    .cost_in_credits
+    .unwrap_or(0);
 
-  let cost = cost_estimate.cost_in_credits.unwrap_or(0);
   info!("v2 estimated cost: {} credits", cost);
 
   // 3. Bill wallet
@@ -336,93 +328,28 @@ async fn execute_v2_pipeline(
 
   let kinovi_character_id_map: Option<HashMap<CharacterToken, String>> =
     resolve_kinovi_character_id_map(
-      request.reference_character_tokens.as_deref(),
-      mysql_connection,
+      request.reference_character_tokens.as_deref(), mysql_connection,
     ).await?;
 
-  // 5. Execute: finalize draft (upload to Kinovi if needed) → send
-  let gen_response = execute_pipeline_v2(
-    draft_or_request,
-    server_state,
-    media_file_urls_as_strings.as_ref(),
-    kinovi_character_id_map.as_ref(),
-    billing.maybe_wallet_ledger_entry_token.as_ref(),
-    mysql_connection,
+  // 5. Execute: finalize draft → send
+  let response = execute_pipeline_v2(
+    draft_or_request, server_state,
+    media_file_urls_as_strings.as_ref(), kinovi_character_id_map.as_ref(),
+    billing.maybe_wallet_ledger_entry_token.as_ref(), mysql_connection,
   ).await?;
 
-  // 6. DB writes
-  let ip_address = get_request_ip(http_request);
-  let mut transaction = mysql_connection.begin().await.map_err(|err| {
-    error!("Error starting MySQL transaction: {:?}", err);
-    AdvancedCommonWebError::from_error(err)
-  })?;
-
-  let prompt_token = insert_prompt_and_context(
-    request, &mut transaction, user_token, maybe_prompt_model_type,
-    None, // TODO: determine generation_mode from v2 response
-    &ip_address,
-  ).await;
-
-  let job_token = match gen_response {
-    GenerateVideoResponse::Seedance2Pro(ref payload) => {
-      insert_seedance2pro_jobs(
-        &payload.order_id,
-        None, // v2 Kinovi returns single order_id per request
-        &billing.apriori_job_token,
-        idempotency_token,
-        prompt_token.as_ref(),
-        billing.maybe_wallet_ledger_entry_token.as_ref(),
-        user_token,
-        maybe_avt_token,
-        &ip_address,
-        &mut transaction,
-      ).await?
-    }
-    GenerateVideoResponse::Artcraft(ref payload) => {
-      // Artcraft recursive: job was already created by the inner omni call
-      payload.inference_job_token.clone()
-    }
-    GenerateVideoResponse::Fal(ref payload) => {
-      let external_id = payload.request_id.clone().unwrap_or_default();
-      insert_fal_job(
-        &external_id,
-        &billing.apriori_job_token,
-        idempotency_token,
-        prompt_token.as_ref(),
-        user_token,
-        maybe_avt_token,
-        &ip_address,
-        &mut transaction,
-      ).await?
-    }
-    _ => {
-      error!("Unexpected v2 response variant: {:?}", gen_response);
-      return Err(AdvancedCommonWebError::server_error_with_message("Unexpected v2 response"));
-    }
-  };
-
-  transaction.commit().await.map_err(|err| {
-    error!("Error committing transaction: {:?}", err);
-    AdvancedCommonWebError::from_error(err)
-  })?;
-
-  Ok(Json(OmniGenVideoGenerateResponse {
-    success: true,
-    inference_job_token: job_token,
-  }))
+  Ok(PipelineResult { billing, response })
 }
 
-// ==================== SHARED HELPERS ==================== //
+// ==================== SHARED DB HELPERS ==================== //
 
-/// Insert prompt and prompt context items. Returns the prompt token if successful.
 async fn insert_prompt_and_context(
   request: &OmniGenVideoCostAndGenerateRequest,
   transaction: &mut sqlx::Transaction<'_, sqlx::MySql>,
   user_token: &tokens::tokens::users::UserToken,
   maybe_prompt_model_type: Option<CommonModelType>,
-  maybe_generation_mode: Option<enums::common::generation::common_generation_mode::CommonGenerationMode>,
   ip_address: &str,
-) -> Option<tokens::tokens::prompts::PromptToken> {
+) {
   let prompt_result = insert_prompt(InsertPromptArgs {
     maybe_apriori_prompt_token: None,
     prompt_type: PromptType::ArtcraftApp,
@@ -432,7 +359,7 @@ async fn insert_prompt_and_context(
     maybe_positive_prompt: request.prompt.as_deref(),
     maybe_negative_prompt: request.negative_prompt.as_deref(),
     maybe_other_args: None,
-    maybe_generation_mode,
+    maybe_generation_mode: None,
     maybe_aspect_ratio: None,
     maybe_resolution: None,
     maybe_batch_count: request.video_batch_count.map(|c| c as u8),
@@ -493,8 +420,6 @@ async fn insert_prompt_and_context(
       }
     }
   }
-
-  prompt_token
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -503,7 +428,6 @@ async fn insert_seedance2pro_jobs(
   maybe_additional_order_ids: Option<&[String]>,
   apriori_job_token: &InferenceJobToken,
   idempotency_token: &str,
-  prompt_token: Option<&tokens::tokens::prompts::PromptToken>,
   maybe_wallet_ledger_entry_token: Option<&tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken>,
   user_token: &tokens::tokens::users::UserToken,
   maybe_avt_token: Option<&tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken>,
@@ -516,17 +440,8 @@ async fn insert_seedance2pro_jobs(
   let mut all_job_tokens: Vec<InferenceJobToken> = Vec::with_capacity(order_ids.len());
 
   for (i, order_id) in order_ids.iter().enumerate() {
-    let job_token = if i == 0 {
-      apriori_job_token.clone()
-    } else {
-      InferenceJobToken::generate()
-    };
-
-    let idempotency_str = if i == 0 {
-      idempotency_token.to_string()
-    } else {
-      format!("{}-batch-{}", idempotency_token, i)
-    };
+    let job_token = if i == 0 { apriori_job_token.clone() } else { InferenceJobToken::generate() };
+    let idempotency_str = if i == 0 { idempotency_token.to_string() } else { format!("{}-batch-{}", idempotency_token, i) };
 
     let db_result = insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token(
       InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs {
@@ -534,7 +449,7 @@ async fn insert_seedance2pro_jobs(
         uuid_idempotency_token: &idempotency_str,
         maybe_external_third_party_id: order_id,
         maybe_inference_args: None,
-        maybe_prompt_token: prompt_token,
+        maybe_prompt_token: None,
         maybe_wallet_ledger_entry_token,
         maybe_creator_user_token: Some(user_token),
         maybe_avt_token,
@@ -549,9 +464,7 @@ async fn insert_seedance2pro_jobs(
       Ok(token) => all_job_tokens.push(token),
       Err(err) => {
         warn!("Error inserting seedance2pro inference job (order_id={}): {:?}", order_id, err);
-        if i == 0 {
-          return Err(AdvancedCommonWebError::from_error(err));
-        }
+        if i == 0 { return Err(AdvancedCommonWebError::from_error(err)); }
       }
     }
   }
@@ -567,13 +480,11 @@ async fn insert_fal_job(
   external_job_id: &str,
   apriori_job_token: &InferenceJobToken,
   idempotency_token: &str,
-  prompt_token: Option<&tokens::tokens::prompts::PromptToken>,
   user_token: &tokens::tokens::users::UserToken,
   maybe_avt_token: Option<&tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken>,
   ip_address: &str,
   transaction: &mut sqlx::Transaction<'_, sqlx::MySql>,
 ) -> Result<InferenceJobToken, AdvancedCommonWebError> {
-  // TODO: Pass maybe_wallet_ledger_entry_token to fal jobs once the fal insert supports it.
   let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(
     InsertGenericInferenceForFalWithAprioriJobTokenArgs {
       apriori_job_token,
@@ -581,7 +492,7 @@ async fn insert_fal_job(
       maybe_external_third_party_id: external_job_id,
       fal_category: FalCategory::VideoGeneration,
       maybe_inference_args: None,
-      maybe_prompt_token: prompt_token,
+      maybe_prompt_token: None,
       maybe_creator_user_token: Some(user_token),
       maybe_avt_token,
       creator_ip_address: ip_address,
@@ -603,8 +514,6 @@ async fn insert_fal_job(
   }
 }
 
-/// Resolve character tokens to a map of CharacterToken → Kinovi character ID.
-/// Used by pipeline v2's draft context.
 async fn resolve_kinovi_character_id_map(
   maybe_tokens: Option<&[CharacterToken]>,
   connection: &mut sqlx::pool::PoolConnection<sqlx::MySql>,
@@ -621,9 +530,7 @@ async fn resolve_kinovi_character_id_map(
 
   let map: HashMap<CharacterToken, String> = characters.iter()
     .filter(|c| c.is_active)
-    .filter_map(|c| {
-      c.kinovi_character_id.as_ref().map(|kid| (c.token.clone(), kid.clone()))
-    })
+    .filter_map(|c| c.kinovi_character_id.as_ref().map(|kid| (c.token.clone(), kid.clone())))
     .collect();
 
   if map.is_empty() { Ok(None) } else { Ok(Some(map)) }