upgrade / refresh tokens

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;

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,
 }

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<SoraJwtBearerToken>,
+  pub (crate) sora_sentinel: Option<SoraSentinel>,
+}
+
+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()),
+    )
+  }
+}

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<Self> {
-    let jwt_claims = lightweight_jwt_parse(&token)?;
+    let jwt_claims = lightweight_sora_jwt_parse(&token)?;
     Ok(SoraJwtBearerToken {
       token,
       jwt_claims,

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,
 }

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;
+pub mod utils;
+mod recipes;

crates/lib/api_clients/openai_sora_client/src/recipes/mod.rs

@@ -0,0 +1 @@
+mod upgrade_or_renew_session;

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(())
+}

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<SoraJwtClaims> {
+pub fn lightweight_sora_jwt_parse(token: &str) -> AnyhowResult<SoraJwtClaims> {
   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<String> {
 
 #[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);

crates/lib/api_clients/openai_sora_client/src/utils/mod.rs

@@ -1 +1 @@
-pub mod lightweight_jwt_parse;
+pub mod lightweight_sora_jwt_parse;