tenant.openagents/omega

1use std::{fmt, sync::Arc};
2
3use async_lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};
4use cloud_api_types::OrganizationId;
5
6use crate::{ClientApiError, CloudApiClient};
7
8#[derive(Clone, Default)]
9pub struct LlmApiToken(Arc<RwLock<Option<CachedLlmApiToken>>>);
10
11struct CachedLlmApiToken {
12    /// The organization ID the token was minted for.
13    organization_id: OrganizationId,
14    token: String,
15}
16
17impl fmt::Debug for CachedLlmApiToken {
18    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
19        f.debug_struct("CachedLlmApiToken")
20            .field("organization_id", &self.organization_id)
21            .field("token", &"<redacted>")
22            .finish()
23    }
24}
25
26impl LlmApiToken {
27    /// Returns the cached LLM token, fetching a fresh one if none has been
28    /// cached yet or if the cached token was minted for a different
29    /// organization. The returned token is not validated; callers must
30    /// be prepared to refresh it (via [`LlmApiToken::refresh`]) if the
31    /// server rejects it.
32    pub async fn cached(
33        &self,
34        client: &CloudApiClient,
35        system_id: Option<String>,
36        organization_id: OrganizationId,
37    ) -> Result<String, ClientApiError> {
38        let lock = self.0.upgradable_read().await;
39        if let Some(CachedLlmApiToken {
40            organization_id: cached_organization_id,
41            token,
42        }) = lock.as_ref()
43            && *cached_organization_id == organization_id
44        {
45            Ok(token.to_string())
46        } else {
47            Self::fetch(
48                RwLockUpgradableReadGuard::upgrade(lock).await,
49                client,
50                system_id,
51                organization_id,
52            )
53            .await
54        }
55    }
56
57    pub async fn refresh(
58        &self,
59        client: &CloudApiClient,
60        system_id: Option<String>,
61        organization_id: OrganizationId,
62    ) -> Result<String, ClientApiError> {
63        Self::fetch(self.0.write().await, client, system_id, organization_id).await
64    }
65
66    pub async fn clear(&self) {
67        *self.0.write().await = None;
68    }
69
70    /// Clears the existing token before attempting to fetch a new one.
71    ///
72    /// Used when switching organizations so that a failed refresh doesn't
73    /// leave a token for the wrong organization.
74    pub async fn clear_and_refresh(
75        &self,
76        client: &CloudApiClient,
77        system_id: Option<String>,
78        organization_id: OrganizationId,
79    ) -> Result<String, ClientApiError> {
80        let mut lock = self.0.write().await;
81        *lock = None;
82        Self::fetch(lock, client, system_id, organization_id).await
83    }
84
85    async fn fetch(
86        mut lock: RwLockWriteGuard<'_, Option<CachedLlmApiToken>>,
87        client: &CloudApiClient,
88        system_id: Option<String>,
89        organization_id: OrganizationId,
90    ) -> Result<String, ClientApiError> {
91        let result = client
92            .create_llm_token(system_id, organization_id.clone())
93            .await;
94        match result {
95            Ok(response) => {
96                *lock = Some(CachedLlmApiToken {
97                    organization_id,
98                    token: response.token.0.clone(),
99                });
100
101                Ok(response.token.0)
102            }
103            Err(err) => {
104                *lock = None;
105                Err(err)
106            }
107        }
108    }
109}
110