137 lines · 4.5 KB · Rust
RawDownload
use super::{Client, UserStore};use anyhow::anyhow;use cloud_api_client::LlmApiToken;use cloud_api_types::websocket_protocol::MessageToClient;use cloud_llm_client::{EXPIRED_LLM_TOKEN_HEADER_NAME, OUTDATED_LLM_TOKEN_HEADER_NAME};use futures::StreamExt;use gpui::{    App, AppContext as _, Context, Entity, EventEmitter, Global, ReadGlobal as _, Subscription,    Task, TaskExt,};use std::sync::Arc; pub trait NeedsLlmTokenRefresh {    /// Returns whether the LLM token needs to be refreshed.    fn needs_llm_token_refresh(&self) -> bool;} impl NeedsLlmTokenRefresh for http_client::Response<http_client::AsyncBody> {    fn needs_llm_token_refresh(&self) -> bool {        self.headers().get(EXPIRED_LLM_TOKEN_HEADER_NAME).is_some()            || self.headers().get(OUTDATED_LLM_TOKEN_HEADER_NAME).is_some()    }} enum TokenRefreshMode {    Refresh,    ClearAndRefresh,} pub fn global_llm_token(cx: &App) -> LlmApiToken {    RefreshLlmTokenListener::global(cx)        .read(cx)        .llm_api_token        .clone()} struct GlobalRefreshLlmTokenListener(Entity<RefreshLlmTokenListener>); impl Global for GlobalRefreshLlmTokenListener {} pub struct LlmTokenRefreshedEvent; pub struct RefreshLlmTokenListener {    client: Arc<Client>,    user_store: Entity<UserStore>,    llm_api_token: LlmApiToken,    _clear_llm_token_on_sign_out: Task<()>,    _subscription: Subscription,} impl EventEmitter<LlmTokenRefreshedEvent> for RefreshLlmTokenListener {} impl RefreshLlmTokenListener {    pub fn register(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {        let listener = cx.new(|cx| RefreshLlmTokenListener::new(client, user_store, cx));        cx.set_global(GlobalRefreshLlmTokenListener(listener));    }     pub fn global(cx: &App) -> Entity<Self> {        GlobalRefreshLlmTokenListener::global(cx).0.clone()    }     fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {        client.add_message_to_client_handler({            let this = cx.weak_entity();            move |message, cx| {                if let Some(this) = this.upgrade() {                    Self::handle_refresh_llm_token(this, message, cx);                }            }        });         let subscription = cx.subscribe(&user_store, |this, _user_store, event, cx| {            if matches!(event, super::user::Event::OrganizationChanged) {                this.refresh(TokenRefreshMode::ClearAndRefresh, cx);            }        });         let llm_api_token = LlmApiToken::default();        let mut status = client.status();        let clear_llm_token_on_sign_out = cx.spawn({            let llm_api_token = llm_api_token.clone();            async move |_this, _cx| {                while let Some(status) = status.next().await {                    if status.is_signed_out() {                        llm_api_token.clear().await;                    }                }            }        });         Self {            client,            user_store,            llm_api_token,            _clear_llm_token_on_sign_out: clear_llm_token_on_sign_out,            _subscription: subscription,        }    }     fn refresh(&self, mode: TokenRefreshMode, cx: &mut Context<Self>) {        let client = self.client.clone();        let llm_api_token = self.llm_api_token.clone();        let organization_id = self            .user_store            .read(cx)            .current_organization()            .map(|organization| organization.id.clone());        cx.spawn(async move |this, cx| {            let organization_id =                organization_id.ok_or_else(|| anyhow!("No organization selected."))?;             match mode {                TokenRefreshMode::Refresh => {                    client                        .refresh_llm_token(&llm_api_token, organization_id)                        .await?;                }                TokenRefreshMode::ClearAndRefresh => {                    client                        .clear_and_refresh_llm_token(&llm_api_token, organization_id)                        .await?;                }            }            this.update(cx, |_this, cx| cx.emit(LlmTokenRefreshedEvent))        })        .detach_and_log_err(cx);    }     fn handle_refresh_llm_token(this: Entity<Self>, message: &MessageToClient, cx: &mut App) {        match message {            MessageToClient::UserUpdated => {                this.update(cx, |this, cx| this.refresh(TokenRefreshMode::Refresh, cx));            }        }    }}