diff --git a/crates/schema/public/enums/src/by_table/users/user_feature_flag.rs b/crates/schema/public/enums/src/by_table/users/user_feature_flag.rs index 05f4632175..7a05ee4206 100644 --- a/crates/schema/public/enums/src/by_table/users/user_feature_flag.rs +++ b/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); } } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/models/video/omni_gen_video_models_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/models/video/omni_gen_video_models_handler.rs index 4642cb1c5d..2ee897799f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/models/video/omni_gen_video_models_handler.rs +++ b/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>, _query: Query, - _server_state: web::Data>, -) -> Result, CommonWebError> { +) -> Result, 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 { + + 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) +} diff --git a/crates/service/web/storyteller_web/src/http_server/session/lookup/user_session_feature_flags.rs b/crates/service/web/storyteller_web/src/http_server/session/lookup/user_session_feature_flags.rs index d31e0d8a56..903a9a9516 100644 --- a/crates/service/web/storyteller_web/src/http_server/session/lookup/user_session_feature_flags.rs +++ b/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)]