tenant.openagents/omega

1use super::{Client, UserStore};
2use anyhow::anyhow;
3use cloud_api_client::LlmApiToken;
4use cloud_api_types::websocket_protocol::MessageToClient;
5use cloud_llm_client::{EXPIRED_LLM_TOKEN_HEADER_NAME, OUTDATED_LLM_TOKEN_HEADER_NAME};
6use futures::StreamExt;
7use gpui::{
8    App, AppContext as _, Context, Entity, EventEmitter, Global, ReadGlobal as _, Subscription,
9    Task, TaskExt,
10};
11use std::sync::Arc;
12
13pub trait NeedsLlmTokenRefresh {
14    /// Returns whether the LLM token needs to be refreshed.
15    fn needs_llm_token_refresh(&self) -> bool;
16}
17
18impl NeedsLlmTokenRefresh for http_client::Response<http_client::AsyncBody> {
19    fn needs_llm_token_refresh(&self) -> bool {
20        self.headers().get(EXPIRED_LLM_TOKEN_HEADER_NAME).is_some()
21            || self.headers().get(OUTDATED_LLM_TOKEN_HEADER_NAME).is_some()
22    }
23}
24
25enum TokenRefreshMode {
26    Refresh,
27    ClearAndRefresh,
28}
29
30pub fn global_llm_token(cx: &App) -> LlmApiToken {
31    RefreshLlmTokenListener::global(cx)
32        .read(cx)
33        .llm_api_token
34        .clone()
35}
36
37struct GlobalRefreshLlmTokenListener(Entity<RefreshLlmTokenListener>);
38
39impl Global for GlobalRefreshLlmTokenListener {}
40
41pub struct LlmTokenRefreshedEvent;
42
43pub struct RefreshLlmTokenListener {
44    client: Arc<Client>,
45    user_store: Entity<UserStore>,
46    llm_api_token: LlmApiToken,
47    _clear_llm_token_on_sign_out: Task<()>,
48    _subscription: Subscription,
49}
50
51impl EventEmitter<LlmTokenRefreshedEvent> for RefreshLlmTokenListener {}
52
53impl RefreshLlmTokenListener {
54    pub fn register(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
55        let listener = cx.new(|cx| RefreshLlmTokenListener::new(client, user_store, cx));
56        cx.set_global(GlobalRefreshLlmTokenListener(listener));
57    }
58
59    pub fn global(cx: &App) -> Entity<Self> {
60        GlobalRefreshLlmTokenListener::global(cx).0.clone()
61    }
62
63    fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
64        client.add_message_to_client_handler({
65            let this = cx.weak_entity();
66            move |message, cx| {
67                if let Some(this) = this.upgrade() {
68                    Self::handle_refresh_llm_token(this, message, cx);
69                }
70            }
71        });
72
73        let subscription = cx.subscribe(&user_store, |this, _user_store, event, cx| {
74            if matches!(event, super::user::Event::OrganizationChanged) {
75                this.refresh(TokenRefreshMode::ClearAndRefresh, cx);
76            }
77        });
78
79        let llm_api_token = LlmApiToken::default();
80        let mut status = client.status();
81        let clear_llm_token_on_sign_out = cx.spawn({
82            let llm_api_token = llm_api_token.clone();
83            async move |_this, _cx| {
84                while let Some(status) = status.next().await {
85                    if status.is_signed_out() {
86                        llm_api_token.clear().await;
87                    }
88                }
89            }
90        });
91
92        Self {
93            client,
94            user_store,
95            llm_api_token,
96            _clear_llm_token_on_sign_out: clear_llm_token_on_sign_out,
97            _subscription: subscription,
98        }
99    }
100
101    fn refresh(&self, mode: TokenRefreshMode, cx: &mut Context<Self>) {
102        let client = self.client.clone();
103        let llm_api_token = self.llm_api_token.clone();
104        let organization_id = self
105            .user_store
106            .read(cx)
107            .current_organization()
108            .map(|organization| organization.id.clone());
109        cx.spawn(async move |this, cx| {
110            let organization_id =
111                organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
112
113            match mode {
114                TokenRefreshMode::Refresh => {
115                    client
116                        .refresh_llm_token(&llm_api_token, organization_id)
117                        .await?;
118                }
119                TokenRefreshMode::ClearAndRefresh => {
120                    client
121                        .clear_and_refresh_llm_token(&llm_api_token, organization_id)
122                        .await?;
123                }
124            }
125            this.update(cx, |_this, cx| cx.emit(LlmTokenRefreshedEvent))
126        })
127        .detach_and_log_err(cx);
128    }
129
130    fn handle_refresh_llm_token(this: Entity<Self>, message: &MessageToClient, cx: &mut App) {
131        match message {
132            MessageToClient::UserUpdated => {
133                this.update(cx, |this, cx| this.refresh(TokenRefreshMode::Refresh, cx));
134            }
135        }
136    }
137}
138