happyhorse feature flags

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

@@ -32,6 +32,14 @@ pub enum UserFeatureFlag {
   /// Access to the faster Seedance 2.0 whitelist
   #[serde(rename = "sd_wl")]
   SeedanceWhitelist,
+
+  /// Access to Happy Horse
+  #[serde(rename = "hh")]
+  HappyHorse,
+  
+  /// Access to Happy Horse (with a rate limit)
+  #[serde(rename = "hh_rl")]
+  HappyHorseRateLimit,
 }
 
 // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests!
@@ -48,6 +56,8 @@ impl UserFeatureFlag {
       Self::Upload3d => "upload_3d",
       Self::VideoStyleTransfer => "video_style_transfer",
       Self::SeedanceWhitelist => "sd_wl",
+      Self::HappyHorse => "hh",
+      Self::HappyHorseRateLimit => "hh_rl",
     }
   }
 
@@ -58,6 +68,8 @@ impl UserFeatureFlag {
       "upload_3d" => Ok(Self::Upload3d),
       "video_style_transfer" => Ok(Self::VideoStyleTransfer),
       "sd_wl" => Ok(Self::SeedanceWhitelist),
+      "hh" => Ok(Self::HappyHorse),
+      "hh_rl" => Ok(Self::HappyHorseRateLimit),
       _ => Err(format!("invalid value: {:?}", value)),
     }
   }
@@ -71,6 +83,8 @@ impl UserFeatureFlag {
       Self::Upload3d,
       Self::VideoStyleTransfer,
       Self::SeedanceWhitelist,
+      Self::HappyHorse,
+      Self::HappyHorseRateLimit,
     ])
   }
 }
@@ -90,6 +104,8 @@ mod tests {
       assert_serialization(UserFeatureFlag::Upload3d, "upload_3d");
       assert_serialization(UserFeatureFlag::VideoStyleTransfer, "video_style_transfer");
       assert_serialization(UserFeatureFlag::SeedanceWhitelist, "sd_wl");
+      assert_serialization(UserFeatureFlag::HappyHorse, "hh");
+      assert_serialization(UserFeatureFlag::HappyHorseRateLimit, "hh_rl");
     }
 
     #[test]
@@ -99,6 +115,8 @@ mod tests {
       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");
+      assert_eq!(UserFeatureFlag::HappyHorse.to_str(), "hh");
+      assert_eq!(UserFeatureFlag::HappyHorseRateLimit.to_str(), "hh_rl");
     }
 
     #[test]
@@ -108,18 +126,22 @@ mod tests {
       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_eq!(UserFeatureFlag::from_str("hh").unwrap(), UserFeatureFlag::HappyHorse);
+      assert_eq!(UserFeatureFlag::from_str("hh_rl").unwrap(), UserFeatureFlag::HappyHorseRateLimit);
       assert!(UserFeatureFlag::from_str("foo").is_err());
     }
 
     #[test]
     fn all_variants() {
       let mut variants = UserFeatureFlag::all_variants();
-      assert_eq!(variants.len(), 5);
+      assert_eq!(variants.len(), 7);
       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(), Some(UserFeatureFlag::HappyHorse));
+      assert_eq!(variants.pop_first(), Some(UserFeatureFlag::HappyHorseRateLimit));
       assert_eq!(variants.pop_first(), None);
     }
   }

crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/models/video/omni_gen_video_models_handler.rs

@@ -5,10 +5,13 @@ use crate::http_server::common_responses::common_web_error::CommonWebError;
 use crate::state::server_state::ServerState;
 use actix_web::web::{Json, Query};
 use actix_web::{web, HttpRequest};
+use log::warn;
 use artcraft_api_defs::omni_gen::models::omni_gen_video_models::{
   OmniGenVideoModelsQuery,
   OmniGenVideoModelsResponse,
 };
+use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
+use crate::http_server::session::lookup::user_session_feature_flags::UserSessionFeatureFlags;
 
 /// List available video models.
 #[utoipa::path(
@@ -22,10 +25,40 @@ use artcraft_api_defs::omni_gen::models::omni_gen_video_models::{
   ),
 )]
 pub async fn omni_gen_video_models_handler(
-  _http_request: HttpRequest,
+  http_request: HttpRequest,
+  server_state: web::Data<Arc<ServerState>>,
   _query: Query<OmniGenVideoModelsQuery>,
-  _server_state: web::Data<Arc<ServerState>>,
-) -> Result<Json<OmniGenVideoModelsResponse>, CommonWebError> {
+) -> Result<Json<OmniGenVideoModelsResponse>, AdvancedCommonWebError> {
   let response = (*OMNI_GEN_VIDEO_MODELS_AND_PROVIDERS).clone();
   Ok(Json(response))
 }
+
+pub async fn can_see_happy_horse(
+  http_request: &HttpRequest,
+  server_state: &ServerState,
+) -> Result<bool, AdvancedCommonWebError> {
+
+  let mut mysql_connection = server_state.mysql_pool.acquire().await?;
+
+  let maybe_user_session = server_state
+      .session_checker
+      .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection)
+      .await
+      .map_err(|e| {
+        warn!("Session checker error: {:?}", e);
+        AdvancedCommonWebError::from(e)
+      })?;
+  
+  let user_session = match maybe_user_session {
+    Some(session) => session,
+    None => return Ok(false),
+  };
+
+  let user_feature_flags =
+      UserSessionFeatureFlags::new(user_session.maybe_feature_flags.as_deref());
+
+  let can_use_happy_horse = user_feature_flags.can_use_happy_horse();
+  let can_use_happy_horse_ratelimited = user_feature_flags.can_use_happy_horse_rate_limited();
+
+  Ok(can_use_happy_horse || can_use_happy_horse_ratelimited)
+}

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

@@ -89,6 +89,14 @@ impl UserSessionFeatureFlags {
   pub fn has_seedance_whitelist(&self) -> bool {
     self.feature_flags.contains(&UserFeatureFlag::SeedanceWhitelist)
   }
+
+  pub fn can_use_happy_horse(&self) -> bool {
+    self.feature_flags.contains(&UserFeatureFlag::SeedanceWhitelist)
+  }
+
+  pub fn can_use_happy_horse_rate_limited(&self) -> bool {
+    self.feature_flags.contains(&UserFeatureFlag::SeedanceWhitelist)
+  }
 }
 
 #[cfg(test)]