cleanup

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

@@ -7,11 +7,14 @@ use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost
 use artcraft_router::api::provider::Provider;
 use artcraft_router::client::router_client::RouterClient;
 use artcraft_router::client::router_seedance2pro_client::RouterSeedance2ProClient;
+use artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse;
 use artcraft_router::generate::generate_video_v2::video_generation_draft_context::VideoGenerationDraftContext;
 use artcraft_router::generate::generate_video_v2::video_generation_draft_or_request::VideoGenerationDraftOrRequest;
 use artcraft_router::generate::generate_video_v2::video_generation_request::VideoGenerationRequest;
 use seedance2pro_client::creds::seedance2pro_session::Seedance2ProSession;
+use tokens::tokens::characters::CharacterToken;
 use tokens::tokens::media_files::MediaFileToken;
+use tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken;
 
 use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
 use crate::http_server::endpoint_helpers::refund_wallet_after_api_failure::refund_wallet_after_api_failure;
@@ -37,6 +40,8 @@ pub async fn run_pipeline_v2(
     AdvancedCommonWebError::from_error(e)
   })?;
 
+  let provider = draft_or_request.get_provider();
+
   // 2. Calculate cost (swap provider to Artcraft for billing)
   let mut cost_builder = hydrate_to_router_request(request)?;
   cost_builder.provider = Provider::Artcraft;
@@ -65,54 +70,24 @@ pub async fn run_pipeline_v2(
       map.iter().map(|(k, v)| (k.clone(), v.to_string())).collect()
     });
 
-  let kinovi_character_id_map = resolve_kinovi_character_ids(
-    request.reference_character_tokens.as_deref(),
-    mysql_connection,
-  ).await?;
+  let kinovi_character_id_map: Option<HashMap<CharacterToken, String>> =
+    resolve_kinovi_character_ids(
+      request.reference_character_tokens.as_deref(),
+      mysql_connection,
+    ).await?;
 
-  // 5. Finalize draft → request (or use request directly)
-  let (video_request, is_kinovi) = match draft_or_request {
-    VideoGenerationDraftOrRequest::Request(request) => {
-      let is_kinovi = matches!(
-        request,
-        VideoGenerationRequest::KinoviSeedance2p0(_) | VideoGenerationRequest::KinoviSeedance2p0Fast(_)
-      );
-      (request, is_kinovi)
-    }
-    VideoGenerationDraftOrRequest::Draft(draft) => {
-      let seedance2pro_session = Seedance2ProSession::from_cookies_string(
-        server_state.seedance2pro.cookies.clone()
-      );
-      let seedance2pro_client = RouterSeedance2ProClient::new(seedance2pro_session);
-      let router_client = RouterClient::Seedance2Pro(seedance2pro_client);
+  // 5. Upload media (if draft) and generate video.
+  //    The entire block is wrapped so Kinovi failures trigger a refund.
+  let result = upload_and_generate(
+    draft_or_request,
+    server_state,
+    media_file_urls_as_strings.as_ref(),
+    kinovi_character_id_map.as_ref(),
+  ).await;
 
-      let draft_context = VideoGenerationDraftContext {
-        client: Some(&router_client),
-        media_file_to_artcraft_url_map: media_file_urls_as_strings.as_ref(),
-        character_token_to_kinovi_id_map: kinovi_character_id_map.as_ref(),
-      };
-
-      let request = draft.finalize(draft_context).await.map_err(|err| {
-        warn!("Failed to finalize v2 draft: {:?}", err);
-        AdvancedCommonWebError::from_error(err)
-      })?;
-
-      // Drafts are always Kinovi (Artcraft never produces drafts).
-      (request, true)
-    }
-  };
-
-  // 6. Build client and send
-  let session = Seedance2ProSession::from_cookies_string(
-    server_state.seedance2pro.cookies.clone()
-  );
-  let client = RouterClient::Seedance2Pro(RouterSeedance2ProClient::new(session));
-
-  let result = video_request.send_request(&client).await;
-
-  // 7. On failure, refund wallet if this is a Kinovi request
+  // 6. On failure, refund wallet for Kinovi requests.
   if let Err(ref err) = result {
-    if is_kinovi {
+    if matches!(provider, Provider::Seedance2Pro) {
       if let Some(ledger_entry_token) = billing.maybe_wallet_ledger_entry_token.as_ref() {
         warn!("Kinovi v2 generation failed, issuing refund for {}: {:?}", ledger_entry_token.as_str(), err);
         if let Err(refund_err) = refund_wallet_after_api_failure(ledger_entry_token, mysql_connection).await {
@@ -122,12 +97,67 @@ pub async fn run_pipeline_v2(
     }
   }
 
-  let response = result.map_err(|err| {
-    warn!("v2 video generation failed: {:?}", err);
-    AdvancedCommonWebError::from_error(err)
-  })?;
-
+  let response = result?;
   info!("v2 generation response: {:?}", response);
 
   Ok(PipelineResult { billing, response })
 }
+
+/// Finalize the draft (uploading media if needed), then send the generation request.
+///
+/// This is the block that gets refunded on failure for Kinovi providers.
+async fn upload_and_generate(
+  draft_or_request: VideoGenerationDraftOrRequest,
+  server_state: &ServerState,
+  media_file_urls_by_token: Option<&HashMap<MediaFileToken, String>>,
+  kinovi_character_ids: Option<&HashMap<CharacterToken, String>>,
+) -> Result<GenerateVideoResponse, AdvancedCommonWebError> {
+  // Finalize draft → request (or use request directly)
+  let video_request = match draft_or_request {
+    VideoGenerationDraftOrRequest::Request(request) => request,
+    VideoGenerationDraftOrRequest::Draft(draft) => {
+      let session = Seedance2ProSession::from_cookies_string(
+        server_state.seedance2pro.cookies.clone()
+      );
+      let seedance2pro_client = RouterSeedance2ProClient::new(session);
+      let router_client = RouterClient::Seedance2Pro(seedance2pro_client);
+
+      let draft_context = VideoGenerationDraftContext {
+        client: Some(&router_client),
+        media_file_to_artcraft_url_map: media_file_urls_by_token,
+        character_token_to_kinovi_id_map: kinovi_character_ids,
+      };
+
+      draft.finalize(draft_context).await.map_err(|err| {
+        warn!("Failed to finalize v2 draft: {:?}", err);
+        AdvancedCommonWebError::from_error(err)
+      })?
+    }
+  };
+
+  // Build client and send
+  let client = build_client(&video_request, server_state);
+
+  video_request.send_request(&client).await.map_err(|err| {
+    warn!("v2 video generation failed: {:?}", err);
+    AdvancedCommonWebError::from_error(err)
+  })
+}
+
+fn build_client(request: &VideoGenerationRequest, server_state: &ServerState) -> RouterClient {
+  match request.get_provider() {
+    Provider::Seedance2Pro => {
+      let session = Seedance2ProSession::from_cookies_string(
+        server_state.seedance2pro.cookies.clone()
+      );
+      RouterClient::Seedance2Pro(RouterSeedance2ProClient::new(session))
+    }
+    // For future providers, add arms here.
+    _ => {
+      let session = Seedance2ProSession::from_cookies_string(
+        server_state.seedance2pro.cookies.clone()
+      );
+      RouterClient::Seedance2Pro(RouterSeedance2ProClient::new(session))
+    }
+  }
+}