seedance2 pro whitelist

_database/sql/migrations_squashed/user_tables/users.sql

@@ -66,6 +66,7 @@ CREATE TABLE users (
   --   - "explore_media"
   --   - "studio"
   --   - "video_style_transfer"
+  --   - "sd_wl" (seedance whitelist)
   maybe_feature_flags VARCHAR(255) DEFAULT NULL,
 
   -- DEPRECATED

crates/schema/public/enums/src/by_table/users/user_feature_flag.rs

@@ -28,6 +28,10 @@ pub enum UserFeatureFlag {
 
   /// Access to video style transfer
   VideoStyleTransfer,
+
+  /// Access to the faster Seedance 2.0 whitelist
+  #[serde(rename = "sd_wl")]
+  SeedanceWhitelist,
 }
 
 // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests!
@@ -43,6 +47,7 @@ impl UserFeatureFlag {
       Self::Studio => "studio",
       Self::Upload3d => "upload_3d",
       Self::VideoStyleTransfer => "video_style_transfer",
+      Self::SeedanceWhitelist => "sd_wl",
     }
   }
 
@@ -52,6 +57,7 @@ impl UserFeatureFlag {
       "studio" => Ok(Self::Studio),
       "upload_3d" => Ok(Self::Upload3d),
       "video_style_transfer" => Ok(Self::VideoStyleTransfer),
+      "sd_wl" => Ok(Self::SeedanceWhitelist),
       _ => Err(format!("invalid value: {:?}", value)),
     }
   }
@@ -64,6 +70,7 @@ impl UserFeatureFlag {
       Self::Studio,
       Self::Upload3d,
       Self::VideoStyleTransfer,
+      Self::SeedanceWhitelist,
     ])
   }
 }
@@ -82,6 +89,7 @@ mod tests {
       assert_serialization(UserFeatureFlag::Studio, "studio");
       assert_serialization(UserFeatureFlag::Upload3d, "upload_3d");
       assert_serialization(UserFeatureFlag::VideoStyleTransfer, "video_style_transfer");
+      assert_serialization(UserFeatureFlag::SeedanceWhitelist, "sd_wl");
     }
 
     #[test]
@@ -90,6 +98,7 @@ mod tests {
       assert_eq!(UserFeatureFlag::Studio.to_str(), "studio");
       assert_eq!(UserFeatureFlag::Upload3d.to_str(), "upload_3d");
       assert_eq!(UserFeatureFlag::VideoStyleTransfer.to_str(), "video_style_transfer");
+      assert_eq!(UserFeatureFlag::SeedanceWhitelist.to_str(), "sd_wl");
     }
 
     #[test]
@@ -98,17 +107,19 @@ mod tests {
       assert_eq!(UserFeatureFlag::from_str("studio").unwrap(), UserFeatureFlag::Studio);
       assert_eq!(UserFeatureFlag::from_str("upload_3d").unwrap(), UserFeatureFlag::Upload3d);
       assert_eq!(UserFeatureFlag::from_str("video_style_transfer").unwrap(), UserFeatureFlag::VideoStyleTransfer);
+      assert_eq!(UserFeatureFlag::from_str("sd_wl").unwrap(), UserFeatureFlag::SeedanceWhitelist);
       assert!(UserFeatureFlag::from_str("foo").is_err());
     }
 
     #[test]
     fn all_variants() {
       let mut variants = UserFeatureFlag::all_variants();
-      assert_eq!(variants.len(), 4);
+      assert_eq!(variants.len(), 5);
       assert_eq!(variants.pop_first(), Some(UserFeatureFlag::ExploreMedia));
       assert_eq!(variants.pop_first(), Some(UserFeatureFlag::Studio));
       assert_eq!(variants.pop_first(), Some(UserFeatureFlag::Upload3d));
       assert_eq!(variants.pop_first(), Some(UserFeatureFlag::VideoStyleTransfer));
+      assert_eq!(variants.pop_first(), Some(UserFeatureFlag::SeedanceWhitelist));
       assert_eq!(variants.pop_first(), None);
     }
   }

crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/multi_function/seedance_2p0_multi_function_video_gen_handler.rs

@@ -35,7 +35,7 @@ use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_it
 use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs};
 use seedance2pro_client::creds::seedance2pro_session::Seedance2ProSession;
 use seedance2pro_client::requests::generate_video::generate_video::{
-  generate_video, BatchCount, GenerateVideoArgs, Resolution,
+  generate_video, BatchCount, GenerateVideoArgs, GenerateVideoResponse, Resolution,
 };
 use seedance2pro_client::requests::prepare_file_upload::prepare_file_upload::{
   prepare_file_upload, PrepareFileUploadArgs,
@@ -46,6 +46,17 @@ use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::media_files::MediaFileToken;
 use url::Url;
 use url_utils::extension::extract_extension_from_url::{extract_extension_from_url, ExtractExtensions};
+use crate::http_server::session::lookup::user_session_feature_flags::UserSessionFeatureFlags;
+
+// ======================== Result of a successful generation ========================
+
+/// Everything the caller needs after a successful upload + generate cycle.
+struct SeedanceGenerationResult {
+  gen_response: GenerateVideoResponse,
+  generation_mode: CommonGenerationMode,
+}
+
+// ======================== Handler ========================
 
 /// Seedance 2.0 Multi-Function video generation (text-to-video, keyframe, and reference).
 #[utoipa::path(
@@ -138,104 +149,13 @@ pub async fn seedance_2p0_multi_function_video_gen_handler(
         CommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string())
       })?;
 
-  // --- Build seedance2pro session ---
-
-  let session = Seedance2ProSession::from_cookies_string(
-    server_state.seedance2pro.cookies.clone()
-  );
-
-  // --- Upload files to seedance2pro CDN ---
-
-  let start_frame_url = match request.start_frame_media_token.as_ref() {
-    None => None,
-    Some(token) => match file_urls_by_token.get(token) {
-      None => return Err(CommonWebError::BadInputWithSimpleMessage("Start frame media not found.".to_string())),
-      Some(url) => Some(upload_to_seedance2pro(&session, url).await?),
-    }
-  };
-
-  let end_frame_url = match request.end_frame_media_token.as_ref() {
-    None => None,
-    Some(token) => match file_urls_by_token.get(token) {
-      None => return Err(CommonWebError::BadInputWithSimpleMessage("End frame media not found.".to_string())),
-      Some(url) => Some(upload_to_seedance2pro(&session, url).await?),
-    }
-  };
-
-  let reference_image_urls = upload_reference_tokens_to_seedance2pro(
-    &session,
-    &file_urls_by_token,
-    request.reference_image_media_tokens.as_deref(),
-    "Reference image",
-  ).await?;
-
-  let reference_video_urls = upload_reference_tokens_to_seedance2pro(
-    &session,
-    &file_urls_by_token,
-    request.reference_video_media_tokens.as_deref(),
-    "Reference video",
-  ).await?;
-
-  let reference_audio_urls = upload_reference_tokens_to_seedance2pro(
-    &session,
-    &file_urls_by_token,
-    request.reference_audio_media_tokens.as_deref(),
-    "Reference audio",
-  ).await?;
-
-  let is_keyframe = request.start_frame_media_token.is_some()
-      || request.end_frame_media_token.is_some();
-
-  let is_reference = request.reference_image_media_tokens.is_some()
-      || request.reference_video_media_tokens.is_some()
-      || request.reference_audio_media_tokens.is_some();
-
-  let generation_mode = if is_keyframe {
-    CommonGenerationMode::Keyframe
-  } else if is_reference {
-    CommonGenerationMode::Reference
-  } else {
-    CommonGenerationMode::Text
-  };
-
-  // --- Map request params to seedance2pro types ---
-
-  let resolution = match request.aspect_ratio {
-    Some(Seedance2p0AspectRatio::Landscape16x9) => Resolution::Landscape16x9,
-    Some(Seedance2p0AspectRatio::Portrait9x16) => Resolution::Portrait9x16,
-    Some(Seedance2p0AspectRatio::Square1x1) => Resolution::Square1x1,
-    Some(Seedance2p0AspectRatio::Standard4x3) => Resolution::Standard4x3,
-    Some(Seedance2p0AspectRatio::Portrait3x4) => Resolution::Portrait3x4,
-    None => Resolution::Landscape16x9,
-  };
-
-  let batch_count = match request.batch_count {
-    Some(Seedance2p0BatchCount::One) | None => BatchCount::One,
-    Some(Seedance2p0BatchCount::Two) => BatchCount::Two,
-    Some(Seedance2p0BatchCount::Four) => BatchCount::Four,
-  };
+  // --- Calculate cost and charge wallet upfront (before uploads) ---
 
+  let resolution = map_resolution(request.aspect_ratio);
+  let batch_count = map_batch_count(request.batch_count);
   let duration_seconds = request.duration_seconds.unwrap_or(5).clamp(4, 15);
-  let prompt = request.prompt.clone().unwrap_or_else(|| "".to_string());
 
-  let video_gen_args = GenerateVideoArgs {
-    session: &session,
-    prompt,
-    resolution,
-    duration_seconds,
-    batch_count,
-    start_frame_url,
-    end_frame_url,
-    reference_image_urls,
-    reference_video_urls,
-    reference_audio_urls,
-    use_face_blur_hack: None,
-    host_override: None,
-  };
-
-  // --- Calculate cost and charge wallet ---
-
-  let cost_in_cents = video_gen_args.estimate_cost_in_usd_cents(); // NB: ArtCraft credits are 1:1 with USD cents.
+  let cost_in_cents = estimate_cost_upfront(resolution, batch_count, duration_seconds);
 
   let apriori_job_token = InferenceJobToken::generate();
 
@@ -248,23 +168,92 @@ pub async fn seedance_2p0_multi_function_video_gen_handler(
     &mut mysql_connection,
   ).await?;
 
-  // --- Call seedance2pro API ---
+  // --- Determine session and generate ---
 
-  // TODO: The refund logic here could use work. We shouldn't hold MySQL transactions open during the HTTP
-  //  transaction (as that could last forever), but this is slightly risky in that a tiny percentage of errors
-  //  may leak. We need to engineer two-phase transactions (a lot of work) to make every step atomic.
+  let user_feature_flags =
+      UserSessionFeatureFlags::new(user_session.maybe_feature_flags.as_deref());
 
-  let gen_response = match generate_video(video_gen_args).await {
-    Ok(resp) => resp,
-    Err(err) => {
-      warn!("Error calling seedance2pro generate_video: {:?}", err);
-      refund_wallet_after_api_failure(&deduction_result.ledger_entry_token, &mut mysql_connection).await?;
+  let is_whitelisted = user_feature_flags.has_seedance_whitelist();
 
-      // TODO: We need to start returning robust failure messages to the caller.
-      return Err(CommonWebError::ServerError);
+  let gen_result = if is_whitelisted {
+    info!("User {:?} is seedance-whitelisted, trying whitelist session first", user_token);
+
+    let whitelist_session = Seedance2ProSession::from_cookies_string(
+      server_state.seedance2pro.cookies_whitelist.clone()
+    );
+
+    let result = upload_and_generate(
+      &whitelist_session,
+      &request,
+      &file_urls_by_token,
+      resolution,
+      batch_count,
+      duration_seconds,
+    ).await;
+
+    match result {
+      Ok(result) => {
+        info!("Whitelist session succeeded for user {:?}", user_token);
+        result
+      }
+      Err(err) => {
+        warn!("Whitelist session failed for user {:?}: {:?}, falling back to regular session", user_token, err);
+
+        let regular_session = Seedance2ProSession::from_cookies_string(
+          server_state.seedance2pro.cookies.clone()
+        );
+
+        let result = upload_and_generate(
+          &regular_session,
+          &request,
+          &file_urls_by_token,
+          resolution,
+          batch_count,
+          duration_seconds,
+        ).await;
+
+        match result {
+          Ok(result) => {
+            info!("Regular session fallback succeeded for whitelisted user {:?}", user_token);
+            result
+          }
+          Err(err) => {
+            warn!("Regular session fallback also failed for user {:?}: {:?}", user_token, err);
+            refund_wallet_after_api_failure(&deduction_result.ledger_entry_token, &mut mysql_connection).await?;
+            return Err(CommonWebError::ServerError);
+          }
+        }
+      }
+    }
+  } else {
+    info!("User {:?} using regular seedance session", user_token);
+
+    let regular_session = Seedance2ProSession::from_cookies_string(
+      server_state.seedance2pro.cookies.clone()
+    );
+
+    let result = upload_and_generate(
+      &regular_session,
+      &request,
+      &file_urls_by_token,
+      resolution,
+      batch_count,
+      duration_seconds,
+    ).await;
+
+    match result {
+      Ok(result) => result,
+      Err(err) => {
+        warn!("Error calling seedance2pro generate_video: {:?}", err);
+        refund_wallet_after_api_failure(&deduction_result.ledger_entry_token, &mut mysql_connection).await?;
+        return Err(CommonWebError::ServerError);
+      }
     }
   };
 
+  let gen_response = gen_result.gen_response;
+  let generation_mode = gen_result.generation_mode;
+
   info!(
     "Seedance2pro task_id={}, order_id={}",
     gen_response.task_id, gen_response.order_id
@@ -438,6 +427,146 @@ pub async fn seedance_2p0_multi_function_video_gen_handler(
   }))
 }
 
+// ======================== Helpers ========================
+
+fn map_resolution(aspect_ratio: Option<Seedance2p0AspectRatio>) -> Resolution {
+  match aspect_ratio {
+    Some(Seedance2p0AspectRatio::Landscape16x9) => Resolution::Landscape16x9,
+    Some(Seedance2p0AspectRatio::Portrait9x16) => Resolution::Portrait9x16,
+    Some(Seedance2p0AspectRatio::Square1x1) => Resolution::Square1x1,
+    Some(Seedance2p0AspectRatio::Standard4x3) => Resolution::Standard4x3,
+    Some(Seedance2p0AspectRatio::Portrait3x4) => Resolution::Portrait3x4,
+    None => Resolution::Landscape16x9,
+  }
+}
+
+fn map_batch_count(batch_count: Option<Seedance2p0BatchCount>) -> BatchCount {
+  match batch_count {
+    Some(Seedance2p0BatchCount::One) | None => BatchCount::One,
+    Some(Seedance2p0BatchCount::Two) => BatchCount::Two,
+    Some(Seedance2p0BatchCount::Four) => BatchCount::Four,
+  }
+}
+
+/// Estimate the cost without needing uploaded URLs. We construct a temporary
+/// `GenerateVideoArgs` with dummy values for the session and URL fields.
+fn estimate_cost_upfront(resolution: Resolution, batch_count: BatchCount, duration_seconds: u8) -> u64 {
+  let dummy_session = Seedance2ProSession::from_cookies_string(String::new());
+  let args = GenerateVideoArgs {
+    session: &dummy_session,
+    prompt: String::new(),
+    resolution,
+    duration_seconds,
+    batch_count,
+    start_frame_url: None,
+    end_frame_url: None,
+    reference_image_urls: None,
+    reference_video_urls: None,
+    reference_audio_urls: None,
+    use_face_blur_hack: None,
+    host_override: None,
+  };
+  args.estimate_cost_in_usd_cents()
+}
+
+/// Uploads all media files and calls generate_video using the given session.
+/// Returns the generation response and the computed generation mode.
+async fn upload_and_generate(
+  session: &Seedance2ProSession,
+  request: &Seedance2p0MultiFunctionVideoGenRequest,
+  file_urls_by_token: &HashMap<MediaFileToken, Url>,
+  resolution: Resolution,
+  batch_count: BatchCount,
+  duration_seconds: u8,
+) -> Result<SeedanceGenerationResult, CommonWebError> {
+
+  // --- Upload files to seedance2pro CDN ---
+
+  let start_frame_url = match request.start_frame_media_token.as_ref() {
+    None => None,
+    Some(token) => match file_urls_by_token.get(token) {
+      None => return Err(CommonWebError::BadInputWithSimpleMessage("Start frame media not found.".to_string())),
+      Some(url) => Some(upload_to_seedance2pro(session, url).await?),
+    }
+  };
+
+  let end_frame_url = match request.end_frame_media_token.as_ref() {
+    None => None,
+    Some(token) => match file_urls_by_token.get(token) {
+      None => return Err(CommonWebError::BadInputWithSimpleMessage("End frame media not found.".to_string())),
+      Some(url) => Some(upload_to_seedance2pro(session, url).await?),
+    }
+  };
+
+  let reference_image_urls = upload_reference_tokens_to_seedance2pro(
+    session,
+    file_urls_by_token,
+    request.reference_image_media_tokens.as_deref(),
+    "Reference image",
+  ).await?;
+
+  let reference_video_urls = upload_reference_tokens_to_seedance2pro(
+    session,
+    file_urls_by_token,
+    request.reference_video_media_tokens.as_deref(),
+    "Reference video",
+  ).await?;
+
+  let reference_audio_urls = upload_reference_tokens_to_seedance2pro(
+    session,
+    file_urls_by_token,
+    request.reference_audio_media_tokens.as_deref(),
+    "Reference audio",
+  ).await?;
+
+  // --- Determine generation mode ---
+
+  let is_keyframe = request.start_frame_media_token.is_some()
+      || request.end_frame_media_token.is_some();
+
+  let is_reference = request.reference_image_media_tokens.is_some()
+      || request.reference_video_media_tokens.is_some()
+      || request.reference_audio_media_tokens.is_some();
+
+  let generation_mode = if is_keyframe {
+    CommonGenerationMode::Keyframe
+  } else if is_reference {
+    CommonGenerationMode::Reference
+  } else {
+    CommonGenerationMode::Text
+  };
+
+  // --- Build args and generate ---
+
+  let prompt = request.prompt.clone().unwrap_or_else(|| "".to_string());
+
+  let video_gen_args = GenerateVideoArgs {
+    session,
+    prompt,
+    resolution,
+    duration_seconds,
+    batch_count,
+    start_frame_url,
+    end_frame_url,
+    reference_image_urls,
+    reference_video_urls,
+    reference_audio_urls,
+    use_face_blur_hack: None,
+    host_override: None,
+  };
+
+  let gen_response = generate_video(video_gen_args).await
+    .map_err(|err| {
+      warn!("Error calling seedance2pro generate_video: {:?}", err);
+      CommonWebError::ServerError
+    })?;
+
+  Ok(SeedanceGenerationResult {
+    gen_response,
+    generation_mode,
+  })
+}
+
 /// Uploads a list of reference media tokens to seedance2pro, returning the resulting URLs.
 /// Returns `None` if the token list is absent or empty.
 async fn upload_reference_tokens_to_seedance2pro(

crates/service/web/storyteller_web/src/http_server/session/lookup/user_session_feature_flags.rs

@@ -83,6 +83,12 @@ impl UserSessionFeatureFlags {
   pub fn clear_flags(&mut self) {
     self.feature_flags = BTreeSet::new();
   }
+  
+  // ----------- Specific Flags ---------- //
+  
+  pub fn has_seedance_whitelist(&self) -> bool {
+    self.feature_flags.contains(&UserFeatureFlag::SeedanceWhitelist)
+  }
 }
 
 #[cfg(test)]

crates/service/web/storyteller_web/src/main.rs

@@ -400,8 +400,6 @@ async fn main() -> AnyhowResult<()> {
   let fal_api_key = FalApiKey::new(easyenv::get_env_string_required("FAL_API_KEY")?);
   let fal_webhook_url = easyenv::get_env_string_required("FAL_WEBHOOK_URL")?;
 
-  let seedance2pro_cookies = easyenv::get_env_string_required("SEEDANCE2PRO_COOKIES")?;
-
   let openai_api_key= easyenv::get_env_string_required("OPENAI_API_KEY")?;
 
   let resend_api_key = easyenv::get_env_string_required("RESEND_API_KEY")?;
@@ -463,7 +461,8 @@ async fn main() -> AnyhowResult<()> {
       webhook_url: fal_webhook_url,
     },
     seedance2pro: Seedance2ProData {
-      cookies: seedance2pro_cookies,
+      cookies: easyenv::get_env_string_required("SEEDANCE2PRO_COOKIES")?,
+      cookies_whitelist: easyenv::get_env_string_required("SEEDANCE2PRO_WHITELIST_COOKIES")?,
     },
     openai: OpenAiData {
       api_key: openai_api_key,

crates/service/web/storyteller_web/src/state/server_state.rs

@@ -317,6 +317,9 @@ pub struct FalData {
 pub struct Seedance2ProData {
   /// Cookie string for seedance2-pro.com
   pub cookies: String,
+
+  /// Cookie string for seedance2-pro.com (Whitelisted account)
+  pub cookies_whitelist: String,
 }
 
 /// OpenAI integration