Conrad Irwin
5b400f541c
Code had been assuming (erroneously, but understandably) that LlmApiToken::acquire would give them a valid token. This is not true, as those tokens expire and you must call refresh explicitly. Add some helpers to do the retry for you, and rename acquire to cached to be clearer about the intent. Closes #ISSUE Release Notes: - Fixed some rare cases where API requests would fail with Unauthorized
79 lines
2.5 KiB
Rust
79 lines
2.5 KiB
Rust
use std::sync::Arc;
|
|
|
|
use async_lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};
|
|
use cloud_api_types::OrganizationId;
|
|
|
|
use crate::{ClientApiError, CloudApiClient};
|
|
|
|
#[derive(Clone, Default)]
|
|
pub struct LlmApiToken(Arc<RwLock<Option<String>>>);
|
|
|
|
impl LlmApiToken {
|
|
/// Returns the cached LLM token, fetching a fresh one only if none has
|
|
/// been cached yet. The returned token is not validated; callers must
|
|
/// be prepared to refresh it (via [`LlmApiToken::refresh`]) if the
|
|
/// server rejects it.
|
|
pub async fn cached(
|
|
&self,
|
|
client: &CloudApiClient,
|
|
system_id: Option<String>,
|
|
organization_id: Option<OrganizationId>,
|
|
) -> Result<String, ClientApiError> {
|
|
let lock = self.0.upgradable_read().await;
|
|
if let Some(token) = lock.as_ref() {
|
|
Ok(token.to_string())
|
|
} else {
|
|
Self::fetch(
|
|
RwLockUpgradableReadGuard::upgrade(lock).await,
|
|
client,
|
|
system_id,
|
|
organization_id,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
|
|
pub async fn refresh(
|
|
&self,
|
|
client: &CloudApiClient,
|
|
system_id: Option<String>,
|
|
organization_id: Option<OrganizationId>,
|
|
) -> Result<String, ClientApiError> {
|
|
Self::fetch(self.0.write().await, client, system_id, organization_id).await
|
|
}
|
|
|
|
/// Clears the existing token before attempting to fetch a new one.
|
|
///
|
|
/// Used when switching organizations so that a failed refresh doesn't
|
|
/// leave a token for the wrong organization.
|
|
pub async fn clear_and_refresh(
|
|
&self,
|
|
client: &CloudApiClient,
|
|
system_id: Option<String>,
|
|
organization_id: Option<OrganizationId>,
|
|
) -> Result<String, ClientApiError> {
|
|
let mut lock = self.0.write().await;
|
|
*lock = None;
|
|
Self::fetch(lock, client, system_id, organization_id).await
|
|
}
|
|
|
|
async fn fetch(
|
|
mut lock: RwLockWriteGuard<'_, Option<String>>,
|
|
client: &CloudApiClient,
|
|
system_id: Option<String>,
|
|
organization_id: Option<OrganizationId>,
|
|
) -> Result<String, ClientApiError> {
|
|
let result = client.create_llm_token(system_id, organization_id).await;
|
|
match result {
|
|
Ok(response) => {
|
|
*lock = Some(response.token.0.clone());
|
|
Ok(response.token.0)
|
|
}
|
|
Err(err) => {
|
|
*lock = None;
|
|
Err(err)
|
|
}
|
|
}
|
|
}
|
|
}
|