diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs b/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs index ce6837da3c..685bd191fd 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/mod.rs @@ -1,3 +1,4 @@ -pub mod sora_jwt_bearer_token; pub mod sora_cookies; +pub mod sora_credential_set; +pub mod sora_jwt_bearer_token; pub mod sora_sentinel; \ No newline at end of file diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs index e8089fb1f4..add1f25fde 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_cookies.rs @@ -1,5 +1,7 @@ /// Cookies are the credential that are always required. /// You can mint a JWT bearer token with just the cookies. + +#[derive(Clone)] pub struct SoraCookies { cookies: String, } diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs new file mode 100644 index 0000000000..ee00f30da9 --- /dev/null +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_credential_set.rs @@ -0,0 +1,26 @@ +use crate::creds::sora_cookies::SoraCookies; +use crate::creds::sora_jwt_bearer_token::SoraJwtBearerToken; +use crate::creds::sora_sentinel::SoraSentinel; + +#[derive(Clone)] +pub struct SoraCredentialSet { + pub (crate) cookies: SoraCookies, + pub (crate) jwt_bearer_token: Option, + pub (crate) sora_sentinel: Option, +} + +impl SoraCredentialSet { + pub fn initialize_with_just_cookies(cookies: SoraCookies) -> Self { + SoraCredentialSet { + cookies, + jwt_bearer_token: None, + sora_sentinel: None, + } + } + + pub fn initialize_with_just_cookies_str(cookies: &str) -> Self { + Self::initialize_with_just_cookies( + SoraCookies::new(cookies.to_string()), + ) + } +} \ No newline at end of file diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs index cb23ea0fab..8826f587b8 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_jwt_bearer_token.rs @@ -1,7 +1,8 @@ -use crate::utils::lightweight_jwt_parse::{lightweight_jwt_parse, SoraJwtClaims}; +use crate::utils::lightweight_sora_jwt_parse::{lightweight_sora_jwt_parse, SoraJwtClaims}; use errors::AnyhowResult; /// Sora bearer tokens are JWT tokens that can be minted with a valid cookie/session. +#[derive(Clone)] pub struct SoraJwtBearerToken { token: String, jwt_claims: SoraJwtClaims, @@ -9,7 +10,7 @@ pub struct SoraJwtBearerToken { impl SoraJwtBearerToken { pub fn new(token: String) -> AnyhowResult { - let jwt_claims = lightweight_jwt_parse(&token)?; + let jwt_claims = lightweight_sora_jwt_parse(&token)?; Ok(SoraJwtBearerToken { token, jwt_claims, diff --git a/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs b/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs index 535e7bce5c..f73d7a0e5f 100644 --- a/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs +++ b/crates/lib/api_clients/openai_sora_client/src/creds/sora_sentinel.rs @@ -1,5 +1,6 @@ /// Sora Sentinel tokens can be generated without cookies or JWTs, and are needed for some API calls. +#[derive(Clone)] pub struct SoraSentinel { sentinel: String, } diff --git a/crates/lib/api_clients/openai_sora_client/src/lib.rs b/crates/lib/api_clients/openai_sora_client/src/lib.rs index 84d743a287..1a9dfe7049 100644 --- a/crates/lib/api_clients/openai_sora_client/src/lib.rs +++ b/crates/lib/api_clients/openai_sora_client/src/lib.rs @@ -2,4 +2,5 @@ pub mod credentials; pub mod creds; pub mod requests; pub mod sora_error; -pub mod utils; \ No newline at end of file +pub mod utils; +mod recipes; \ No newline at end of file diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs new file mode 100644 index 0000000000..aa7bd549b8 --- /dev/null +++ b/crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs @@ -0,0 +1 @@ +mod upgrade_or_renew_session; \ No newline at end of file diff --git a/crates/lib/api_clients/openai_sora_client/src/recipes/upgrade_or_renew_session.rs b/crates/lib/api_clients/openai_sora_client/src/recipes/upgrade_or_renew_session.rs new file mode 100644 index 0000000000..38e1cc1af3 --- /dev/null +++ b/crates/lib/api_clients/openai_sora_client/src/recipes/upgrade_or_renew_session.rs @@ -0,0 +1,51 @@ +use crate::creds::sora_credential_set::SoraCredentialSet; +use crate::creds::sora_jwt_bearer_token::SoraJwtBearerToken; +use crate::creds::sora_sentinel::SoraSentinel; +use crate::requests::bearer::generate::generate_bearer_with_cookie; +use crate::requests::sentinel_refresh::generate::token::generate_token; +use chrono::{DateTime, TimeDelta, Utc}; +use errors::AnyhowResult; +use log::info; +use std::ops::Sub; + +const EXPIRATION_DEADLINE : TimeDelta = TimeDelta::hours(12); + +/// Call this at session startup to generate the tokens for the first time. +/// Call this periodically to refresh the JWT as it nears its expiration date. +/// This will not refresh the sentinel token once it's been generated - we can use other flows to accomplish that. +pub async fn upgrade_or_renew_session(sora_credentials: &mut SoraCredentialSet) -> AnyhowResult<()> { + let mut refresh_jwt = !sora_credentials.jwt_bearer_token.is_some(); + let refresh_sentinel = !sora_credentials.sora_sentinel.is_some(); + + if let Some(jwt) = &sora_credentials.jwt_bearer_token { + let now = Utc::now(); + let refresh_deadline = now.sub(EXPIRATION_DEADLINE); + if jwt.jwt_claims().expiration.lt(&now) { + info!("JWT expiration time has elapsed."); + refresh_jwt = true; + } else if jwt.jwt_claims().expiration.lt(&refresh_deadline) { + info!("JWT expiration time is under the expiration deadline."); + refresh_jwt = true; + } + } + + if refresh_jwt { + let cookies = sora_credentials.cookies.get_cookies(); + + info!("Generating new JWT bearer token..."); + let token = generate_bearer_with_cookie(cookies).await?; + + info!("Parsing JWT bearer token."); + let token = SoraJwtBearerToken::new(token)?; + sora_credentials.jwt_bearer_token = Some(token); + } + + if refresh_sentinel { + info!("Generating new sentinel..."); + let token = generate_token().await?; + let token = SoraSentinel::new(token); + sora_credentials.sora_sentinel = Some(token); + } + + Ok(()) +} diff --git a/crates/lib/api_clients/openai_sora_client/src/utils/lightweight_jwt_parse.rs b/crates/lib/api_clients/openai_sora_client/src/utils/lightweight_sora_jwt_parse.rs similarity index 95% rename from crates/lib/api_clients/openai_sora_client/src/utils/lightweight_jwt_parse.rs rename to crates/lib/api_clients/openai_sora_client/src/utils/lightweight_sora_jwt_parse.rs index e06c1dae8f..a1605e82cb 100644 --- a/crates/lib/api_clients/openai_sora_client/src/utils/lightweight_jwt_parse.rs +++ b/crates/lib/api_clients/openai_sora_client/src/utils/lightweight_sora_jwt_parse.rs @@ -19,7 +19,7 @@ pub struct SoraJwtClaims { pub email_verified: bool, } -pub fn lightweight_jwt_parse(token: &str) -> AnyhowResult { +pub fn lightweight_sora_jwt_parse(token: &str) -> AnyhowResult { let [_header_str, claims_str, _signature_str] = split_components(token)?; let claims = decode_base64(claims_str)?; @@ -100,7 +100,7 @@ fn decode_base64(raw: &str) -> AnyhowResult { #[cfg(test)] mod tests { - use crate::utils::lightweight_jwt_parse::lightweight_jwt_parse; + use crate::utils::lightweight_sora_jwt_parse::lightweight_sora_jwt_parse; #[test] fn test_sora_jwt_bearer_token() { @@ -131,7 +131,7 @@ mod tests { // - "nbf": 1744936471, - not before, don't use before date // - "exp": 1745800472, - in 5 days, expiry - let result = lightweight_jwt_parse(token).expect("claim should parse"); + let result = lightweight_sora_jwt_parse(token).expect("claim should parse"); assert_eq!(result.email, "vocodes2020@gmail.com"); assert_eq!(result.email_verified, true); diff --git a/crates/lib/api_clients/openai_sora_client/src/utils/mod.rs b/crates/lib/api_clients/openai_sora_client/src/utils/mod.rs index 8862cb61fa..d8ecdb640c 100644 --- a/crates/lib/api_clients/openai_sora_client/src/utils/mod.rs +++ b/crates/lib/api_clients/openai_sora_client/src/utils/mod.rs @@ -1 +1 @@ -pub mod lightweight_jwt_parse; \ No newline at end of file +pub mod lightweight_sora_jwt_parse; \ No newline at end of file