Skip to repository content1590 lines · 54.9 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T05:55:17.248Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
openai_subscribed.rs
1use anyhow::{Context as _, Result, anyhow};
2use base64::Engine as _;
3use base64::engine::general_purpose::URL_SAFE_NO_PAD;
4use credentials_provider::CredentialsProvider;
5use futures::{FutureExt, StreamExt, future::BoxFuture, future::Shared};
6use gpui::{App, AsyncApp, Context, Entity, SharedString, Task, Window};
7use http_client::{
8 AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest,
9 http::{HeaderName, HeaderValue},
10};
11use language_model::{
12 AuthenticateError, FastModeConfirmation, IconOrSvg, InlineDescription, LanguageModel,
13 LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel,
14 LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId,
15 LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest,
16 LanguageModelToolChoice, ProviderSettingsView, RateLimiter,
17};
18use open_ai::{ReasoningEffort, responses::stream_response};
19use rand::RngCore as _;
20use serde::{Deserialize, Serialize};
21use sha2::{Digest, Sha256};
22use std::sync::Arc;
23use std::time::{SystemTime, UNIX_EPOCH};
24use ui::{ConfiguredApiCard, prelude::*};
25use url::form_urlencoded;
26use util::ResultExt as _;
27
28use crate::provider::open_ai::{OpenAiResponseEventMapper, into_open_ai_response};
29
30const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openai-subscribed");
31const PROVIDER_NAME: LanguageModelProviderName =
32 LanguageModelProviderName::new("ChatGPT Subscription");
33
34const SUBSCRIPTION_DESCRIPTION: &str =
35 "Sign in with your ChatGPT Plus or Pro subscription to use OpenAI models in Omega's built-in agent.";
36
37const CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex";
38const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
39const OPENAI_AUTHORIZE_URL: &str = "https://auth.openai.com/oauth/authorize";
40const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
41
42const CREDENTIALS_KEY: &str = "https://chatgpt.com/backend-api/codex";
43const TOKEN_REFRESH_BUFFER_MS: u64 = 5 * 60 * 1000;
44
45#[derive(Serialize, Deserialize, Clone, Debug)]
46struct CodexCredentials {
47 access_token: String,
48 refresh_token: String,
49 expires_at_ms: u64,
50 account_id: Option<String>,
51 email: Option<String>,
52}
53
54impl CodexCredentials {
55 fn is_expired(&self) -> bool {
56 let now = now_ms();
57 now + TOKEN_REFRESH_BUFFER_MS >= self.expires_at_ms
58 }
59}
60
61pub struct State {
62 credentials: Option<CodexCredentials>,
63 sign_in_task: Option<Task<Result<()>>>,
64 refresh_task: Option<Shared<Task<Result<CodexCredentials, Arc<anyhow::Error>>>>>,
65 load_task: Option<Shared<Task<Result<(), Arc<anyhow::Error>>>>>,
66 credentials_provider: Arc<dyn CredentialsProvider>,
67 auth_generation: u64,
68 last_auth_error: Option<SharedString>,
69}
70
71#[derive(Debug)]
72enum RefreshError {
73 Fatal(anyhow::Error),
74 Transient(anyhow::Error),
75}
76
77impl std::fmt::Display for RefreshError {
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 match self {
80 RefreshError::Fatal(e) => write!(f, "{e}"),
81 RefreshError::Transient(e) => write!(f, "{e}"),
82 }
83 }
84}
85
86impl State {
87 fn is_authenticated(&self) -> bool {
88 self.credentials.is_some()
89 }
90
91 fn email(&self) -> Option<&str> {
92 self.credentials.as_ref().and_then(|c| c.email.as_deref())
93 }
94
95 fn is_signing_in(&self) -> bool {
96 self.sign_in_task.is_some()
97 }
98}
99
100pub struct OpenAiSubscribedProvider {
101 http_client: Arc<dyn HttpClient>,
102 state: Entity<State>,
103}
104
105impl OpenAiSubscribedProvider {
106 pub fn new(
107 http_client: Arc<dyn HttpClient>,
108 credentials_provider: Arc<dyn CredentialsProvider>,
109 cx: &mut App,
110 ) -> Self {
111 let state = cx.new(|_cx| State {
112 credentials: None,
113 sign_in_task: None,
114 refresh_task: None,
115 load_task: None,
116 credentials_provider,
117 auth_generation: 0,
118 last_auth_error: None,
119 });
120
121 let provider = Self { http_client, state };
122
123 provider.load_credentials(cx);
124
125 provider
126 }
127
128 fn load_credentials(&self, cx: &mut App) {
129 let state = self.state.downgrade();
130 let load_task = cx
131 .spawn(async move |cx| {
132 let credentials_provider =
133 state.read_with(&*cx, |s, _| s.credentials_provider.clone())?;
134 let result = credentials_provider
135 .read_credentials(CREDENTIALS_KEY, &*cx)
136 .await;
137 state.update(cx, |s, cx| {
138 if let Ok(Some((_, bytes))) = result {
139 match serde_json::from_slice::<CodexCredentials>(&bytes) {
140 Ok(creds) => s.credentials = Some(creds),
141 Err(err) => {
142 log::warn!(
143 "Failed to deserialize ChatGPT subscription credentials: {err}"
144 );
145 }
146 }
147 }
148 s.load_task = None;
149 cx.notify();
150 })?;
151 Ok::<(), Arc<anyhow::Error>>(())
152 })
153 .shared();
154
155 self.state.update(cx, |s, _| {
156 s.load_task = Some(load_task);
157 });
158 }
159
160 fn create_language_model(&self, model: ChatGptModel) -> Arc<dyn LanguageModel> {
161 Arc::new(OpenAiSubscribedLanguageModel {
162 id: LanguageModelId::from(model.id().to_string()),
163 model,
164 state: self.state.clone(),
165 http_client: self.http_client.clone(),
166 request_limiter: RateLimiter::new(4),
167 })
168 }
169}
170
171impl LanguageModelProviderState for OpenAiSubscribedProvider {
172 type ObservableEntity = State;
173
174 fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
175 Some(self.state.clone())
176 }
177}
178
179impl LanguageModelProvider for OpenAiSubscribedProvider {
180 fn id(&self) -> LanguageModelProviderId {
181 PROVIDER_ID
182 }
183
184 fn name(&self) -> LanguageModelProviderName {
185 PROVIDER_NAME
186 }
187
188 fn icon(&self) -> IconOrSvg {
189 IconOrSvg::Icon(IconName::AiOpenAi)
190 }
191
192 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
193 Some(self.create_language_model(ChatGptModel::Gpt55))
194 }
195
196 fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
197 Some(self.create_language_model(ChatGptModel::Gpt56Luna))
198 }
199
200 fn provided_models(&self, _cx: &App) -> Vec<Arc<dyn LanguageModel>> {
201 ChatGptModel::all()
202 .into_iter()
203 .map(|m| self.create_language_model(m))
204 .collect()
205 }
206
207 fn is_authenticated(&self, cx: &App) -> bool {
208 self.state.read(cx).is_authenticated()
209 }
210
211 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
212 if self.is_authenticated(cx) {
213 return Task::ready(Ok(()));
214 }
215 let load_task = self.state.read(cx).load_task.clone();
216 if let Some(load_task) = load_task {
217 let weak_state = self.state.downgrade();
218 cx.spawn(async move |cx| {
219 let _ = load_task.await;
220 let is_auth = weak_state
221 .read_with(&*cx, |s, _| s.is_authenticated())
222 .unwrap_or(false);
223 if is_auth {
224 Ok(())
225 } else {
226 Err(anyhow!(
227 "Sign in with your ChatGPT Plus or Pro subscription to use this provider."
228 )
229 .into())
230 }
231 })
232 } else {
233 Task::ready(Err(anyhow!(
234 "Sign in with your ChatGPT Plus or Pro subscription to use this provider."
235 )
236 .into()))
237 }
238 }
239
240 fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView> {
241 let is_authenticated = self.state.read(cx).is_authenticated();
242 let title = if is_authenticated {
243 None
244 } else {
245 Some("Configure ChatGPT".into())
246 };
247 let description = if is_authenticated {
248 None
249 } else {
250 Some(InlineDescription::Text(SUBSCRIPTION_DESCRIPTION.into()))
251 };
252
253 Some(ProviderSettingsView::Inline(
254 language_model::InlineProviderSettings {
255 title,
256 description,
257 create_view: Arc::new({
258 let state = self.state.clone();
259 let http_client = self.http_client.clone();
260 move |_window, cx| {
261 cx.new(|_cx| ConfigurationView {
262 state: state.clone(),
263 http_client: http_client.clone(),
264 compact: true,
265 })
266 .into()
267 }
268 }),
269 },
270 ))
271 }
272
273 fn authentication_error_message(&self) -> SharedString {
274 "Your ChatGPT subscription session is invalid or has expired. \
275 Sign in again via Settings > AI > LLM Providers to continue."
276 .into()
277 }
278
279 fn missing_credentials_error_message(&self) -> SharedString {
280 "You are not signed in to your ChatGPT account. \
281 Sign in via Settings > AI > LLM Providers to continue."
282 .into()
283 }
284
285 fn fast_mode_confirmation(&self, _cx: &App) -> Option<FastModeConfirmation> {
286 Some(FastModeConfirmation {
287 title: "Enable Fast Mode for OpenAI?".into(),
288 message: "Fast mode sends requests using OpenAI's Priority processing tier, which \
289 targets significantly lower latency than the standard tier and is billed at a \
290 premium per-token rate."
291 .into(),
292 })
293 }
294}
295
296//
297// The ChatGPT Subscription provider routes requests to chatgpt.com/backend-api/codex,
298// which only supports a subset of OpenAI models. This list is maintained separately
299// from the standard OpenAI API model list (open_ai::Model).
300//
301// TODO: The Codex CLI fetches this list dynamically from
302// `GET <codex_base_url>/models?client_version=...` (see
303// codex-rs/codex-api/src/endpoint/models.rs in openai/codex) and falls back to
304// a bundled models.json. Beyond going stale, the static approach also can't
305// model per-account access (e.g. free accounts cannot use gpt-5.4 even though
306// paid accounts can), so the backend still rejects some requests. The bundled
307// list at
308// codex-rs/models-manager/models.json (openai/codex) is the closest
309// approximation; the entries below mirror that file's picker-visible models.
310#[derive(Clone, Debug, PartialEq)]
311enum ChatGptModel {
312 Gpt56Sol,
313 Gpt56Terra,
314 Gpt56Luna,
315 Gpt55,
316 Gpt54,
317 Gpt54Mini,
318}
319
320impl ChatGptModel {
321 fn all() -> Vec<Self> {
322 vec![
323 Self::Gpt56Sol,
324 Self::Gpt56Terra,
325 Self::Gpt56Luna,
326 Self::Gpt55,
327 Self::Gpt54,
328 Self::Gpt54Mini,
329 ]
330 }
331
332 fn id(&self) -> &str {
333 match self {
334 Self::Gpt56Sol => "gpt-5.6-sol",
335 Self::Gpt56Terra => "gpt-5.6-terra",
336 Self::Gpt56Luna => "gpt-5.6-luna",
337 Self::Gpt55 => "gpt-5.5",
338 Self::Gpt54 => "gpt-5.4",
339 Self::Gpt54Mini => "gpt-5.4-mini",
340 }
341 }
342
343 fn display_name(&self) -> &str {
344 match self {
345 Self::Gpt56Sol => "GPT-5.6 Sol",
346 Self::Gpt56Terra => "GPT-5.6 Terra",
347 Self::Gpt56Luna => "GPT-5.6 Luna",
348 Self::Gpt55 => "GPT-5.5",
349 Self::Gpt54 => "GPT-5.4",
350 Self::Gpt54Mini => "GPT-5.4 Mini",
351 }
352 }
353
354 fn max_token_count(&self) -> u64 {
355 match self {
356 Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna => 372_000,
357 Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => 272_000,
358 }
359 }
360
361 fn max_output_tokens(&self) -> Option<u64> {
362 // Codex model metadata does not expose a max output token cap for these
363 // models. Source: openai/codex models-manager/models.json.
364 None
365 }
366
367 fn supports_images(&self) -> bool {
368 true
369 }
370
371 fn default_reasoning_effort(&self) -> Option<ReasoningEffort> {
372 match self {
373 Self::Gpt56Sol => Some(ReasoningEffort::Low),
374 Self::Gpt56Terra | Self::Gpt56Luna | Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => {
375 Some(ReasoningEffort::Medium)
376 }
377 }
378 }
379
380 fn supported_reasoning_efforts(&self) -> &'static [ReasoningEffort] {
381 match self {
382 Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna => &[
383 ReasoningEffort::Low,
384 ReasoningEffort::Medium,
385 ReasoningEffort::High,
386 ReasoningEffort::XHigh,
387 ReasoningEffort::Max,
388 ],
389 Self::Gpt55 | Self::Gpt54 | Self::Gpt54Mini => &[
390 ReasoningEffort::Low,
391 ReasoningEffort::Medium,
392 ReasoningEffort::High,
393 ReasoningEffort::XHigh,
394 ],
395 }
396 }
397
398 fn supports_parallel_tool_calls(&self) -> bool {
399 true
400 }
401
402 fn supports_prompt_cache_key(&self) -> bool {
403 true
404 }
405
406 fn supports_priority(&self) -> bool {
407 match self {
408 Self::Gpt56Sol | Self::Gpt56Terra | Self::Gpt56Luna | Self::Gpt55 | Self::Gpt54 => true,
409 Self::Gpt54Mini => false,
410 }
411 }
412}
413
414struct OpenAiSubscribedLanguageModel {
415 id: LanguageModelId,
416 model: ChatGptModel,
417 state: Entity<State>,
418 http_client: Arc<dyn HttpClient>,
419 request_limiter: RateLimiter,
420}
421
422impl LanguageModel for OpenAiSubscribedLanguageModel {
423 fn id(&self) -> LanguageModelId {
424 self.id.clone()
425 }
426
427 fn name(&self) -> LanguageModelName {
428 LanguageModelName::from(self.model.display_name().to_string())
429 }
430
431 fn provider_id(&self) -> LanguageModelProviderId {
432 PROVIDER_ID
433 }
434
435 fn provider_name(&self) -> LanguageModelProviderName {
436 PROVIDER_NAME
437 }
438
439 fn supports_tools(&self) -> bool {
440 true
441 }
442
443 fn supports_images(&self) -> bool {
444 self.model.supports_images()
445 }
446
447 fn supports_tool_choice(&self, _choice: LanguageModelToolChoice) -> bool {
448 true
449 }
450
451 fn supports_streaming_tools(&self) -> bool {
452 true
453 }
454
455 fn supports_thinking(&self) -> bool {
456 true
457 }
458
459 fn supports_fast_mode(&self) -> bool {
460 self.model.supports_priority()
461 }
462
463 fn supported_effort_levels(&self) -> Vec<LanguageModelEffortLevel> {
464 let default_effort = self.model.default_reasoning_effort();
465 self.model
466 .supported_reasoning_efforts()
467 .iter()
468 .copied()
469 .filter_map(|effort| {
470 let (name, value) = match effort {
471 ReasoningEffort::None => return None,
472 ReasoningEffort::Minimal => ("Minimal", "minimal"),
473 ReasoningEffort::Low => ("Low", "low"),
474 ReasoningEffort::Medium => ("Medium", "medium"),
475 ReasoningEffort::High => ("High", "high"),
476 ReasoningEffort::XHigh => ("Extra High", "xhigh"),
477 ReasoningEffort::Max => ("Max", "max"),
478 };
479
480 Some(LanguageModelEffortLevel {
481 name: name.into(),
482 value: value.into(),
483 is_default: Some(effort) == default_effort,
484 })
485 })
486 .collect()
487 }
488
489 fn telemetry_id(&self) -> String {
490 format!("openai-subscribed/{}", self.model.id())
491 }
492
493 fn max_token_count(&self) -> u64 {
494 self.model.max_token_count()
495 }
496
497 fn max_output_tokens(&self) -> Option<u64> {
498 self.model.max_output_tokens()
499 }
500
501 fn stream_completion(
502 &self,
503 mut request: LanguageModelRequest,
504 cx: &AsyncApp,
505 ) -> BoxFuture<
506 'static,
507 Result<
508 futures::stream::BoxStream<
509 'static,
510 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
511 >,
512 LanguageModelCompletionError,
513 >,
514 > {
515 if !self.model.supports_priority() {
516 request.speed = None;
517 }
518
519 // The Codex backend rejects `max_output_tokens` (`Unsupported parameter`),
520 // unlike the public OpenAI Responses API. Pass `None` so the field is
521 // omitted from the serialized request body entirely.
522 let mut responses_request = match into_open_ai_response(
523 request,
524 self.model.id(),
525 self.model.supports_parallel_tool_calls(),
526 self.model.supports_prompt_cache_key(),
527 /*max_output_tokens*/ None,
528 self.model.default_reasoning_effort(),
529 self.model
530 .supported_reasoning_efforts()
531 .contains(&ReasoningEffort::None),
532 &PROVIDER_ID,
533 ) {
534 Ok(request) => request,
535 Err(error) => return async move { Err(error.into()) }.boxed(),
536 };
537 responses_request.store = Some(false);
538
539 // `into_open_ai_response` already hoists system messages into
540 // `instructions`, which is the only form the Codex backend accepts.
541 // Codex has only ever been sent requests with the field present
542 // (possibly empty), so keep sending it even without system messages.
543 responses_request.instructions.get_or_insert_default();
544
545 let state = self.state.downgrade();
546 let http_client = self.http_client.clone();
547 let request_limiter = self.request_limiter.clone();
548
549 let future = cx.spawn(async move |cx| {
550 let creds = get_fresh_credentials(&state, &http_client, cx).await?;
551
552 let mut header_pairs: Vec<(HeaderName, HeaderValue)> = vec![
553 (
554 HeaderName::from_static("originator"),
555 HeaderValue::from_static("zed"),
556 ),
557 (
558 HeaderName::from_static("openai-beta"),
559 HeaderValue::from_static("responses=experimental"),
560 ),
561 ];
562 if let Some(ref id) = creds.account_id {
563 if !id.is_empty() {
564 if let Ok(value) = HeaderValue::from_str(id) {
565 header_pairs.push((HeaderName::from_static("chatgpt-account-id"), value));
566 }
567 }
568 }
569 let extra_headers = CustomHeaders::new(header_pairs);
570
571 let access_token = creds.access_token.clone();
572 request_limiter
573 .stream(async move {
574 stream_response(
575 http_client.as_ref(),
576 PROVIDER_NAME.0.as_str(),
577 CODEX_BASE_URL,
578 &access_token,
579 responses_request,
580 &extra_headers,
581 )
582 .await
583 .map_err(LanguageModelCompletionError::from)
584 })
585 .await
586 });
587
588 async move {
589 let mapper = OpenAiResponseEventMapper::new(PROVIDER_ID);
590 Ok(mapper.map_stream(future.await?.boxed()).boxed())
591 }
592 .boxed()
593 }
594}
595
596async fn get_fresh_credentials(
597 state: &gpui::WeakEntity<State>,
598 http_client: &Arc<dyn HttpClient>,
599 cx: &mut AsyncApp,
600) -> Result<CodexCredentials, LanguageModelCompletionError> {
601 let (creds, existing_task) = state
602 .read_with(&*cx, |s, _| (s.credentials.clone(), s.refresh_task.clone()))
603 .map_err(LanguageModelCompletionError::Other)?;
604
605 let creds = creds.ok_or(LanguageModelCompletionError::NoApiKey {
606 provider: PROVIDER_NAME,
607 })?;
608
609 if !creds.is_expired() {
610 return Ok(creds);
611 }
612
613 // If another caller is already refreshing, await their result.
614 if let Some(shared_task) = existing_task {
615 return shared_task
616 .await
617 .map_err(|e| LanguageModelCompletionError::Other(anyhow::anyhow!("{e}")));
618 }
619
620 // We are the first caller to notice expiry — spawn the refresh task.
621 let http_client_clone = http_client.clone();
622 let state_clone = state.clone();
623 let refresh_token_value = creds.refresh_token.clone();
624
625 // Capture the generation so we can detect sign-outs that happened during refresh.
626 let generation = state
627 .read_with(&*cx, |s, _| s.auth_generation)
628 .map_err(LanguageModelCompletionError::Other)?;
629
630 let shared_task = cx
631 .spawn(async move |cx| {
632 let result = refresh_token(&http_client_clone, &refresh_token_value).await;
633
634 match result {
635 Ok(refreshed) => {
636 let persist_result: Result<CodexCredentials, Arc<anyhow::Error>> = async {
637 // Check if auth_generation changed (sign-out during refresh).
638 let current_generation = state_clone
639 .read_with(&*cx, |s, _| s.auth_generation)
640 .map_err(|e| Arc::new(e))?;
641 if current_generation != generation {
642 return Err(Arc::new(anyhow!(
643 "Sign-out occurred during token refresh"
644 )));
645 }
646
647 let credentials_provider = state_clone
648 .read_with(&*cx, |s, _| s.credentials_provider.clone())
649 .map_err(|e| Arc::new(e))?;
650
651 let json =
652 serde_json::to_vec(&refreshed).map_err(|e| Arc::new(e.into()))?;
653
654 credentials_provider
655 .write_credentials(CREDENTIALS_KEY, "Bearer", &json, &*cx)
656 .await
657 .map_err(|e| Arc::new(e))?;
658
659 state_clone
660 .update(cx, |s, _| {
661 s.credentials = Some(refreshed.clone());
662 s.refresh_task = None;
663 })
664 .map_err(|e| Arc::new(e))?;
665
666 Ok(refreshed)
667 }
668 .await;
669
670 // Clear refresh_task on failure too.
671 if persist_result.is_err() {
672 let _ = state_clone.update(cx, |s, _| {
673 s.refresh_task = None;
674 });
675 }
676
677 persist_result
678 }
679 Err(RefreshError::Fatal(e)) => {
680 log::error!("ChatGPT subscription token refresh failed fatally: {e:?}");
681 let _ = state_clone.update(cx, |s, cx| {
682 s.refresh_task = None;
683 s.credentials = None;
684 s.last_auth_error =
685 Some("Your session has expired. Please sign in again.".into());
686 cx.notify();
687 });
688 // Also clear the keychain so stale credentials aren't loaded next time.
689 if let Ok(credentials_provider) =
690 state_clone.read_with(&*cx, |s, _| s.credentials_provider.clone())
691 {
692 credentials_provider
693 .delete_credentials(CREDENTIALS_KEY, &*cx)
694 .await
695 .log_err();
696 }
697 Err(Arc::new(e))
698 }
699 Err(RefreshError::Transient(e)) => {
700 log::warn!("ChatGPT subscription token refresh failed transiently: {e:?}");
701 let _ = state_clone.update(cx, |s, _| {
702 s.refresh_task = None;
703 });
704 Err(Arc::new(e))
705 }
706 }
707 })
708 .shared();
709
710 // Store the shared task so concurrent callers can join on it.
711 state
712 .update(cx, |s, _| {
713 s.refresh_task = Some(shared_task.clone());
714 })
715 .map_err(LanguageModelCompletionError::Other)?;
716
717 shared_task
718 .await
719 .map_err(|e| LanguageModelCompletionError::Other(anyhow::anyhow!("{e}")))
720}
721
722#[derive(Deserialize)]
723struct TokenResponse {
724 access_token: String,
725 refresh_token: String,
726 #[serde(default)]
727 id_token: Option<String>,
728 expires_in: u64,
729 #[serde(default)]
730 email: Option<String>,
731}
732
733// The OAuth client registered for `CLIENT_ID` (the Codex CLI's client) only allows
734// `http://localhost:1455/auth/callback` and `http://localhost:1457/auth/callback`
735// as redirect URIs; using anything else (different host, port, or path) causes
736// auth.openai.com to reject the authorize request with a generic `unknown_error`
737// before redirecting back. Keep these in sync with the Codex CLI's redirect URI
738// allow-list (see codex-rs/login/src/server.rs in openai/codex).
739const CODEX_CALLBACK_HOST: &str = "localhost";
740const CODEX_CALLBACK_PORT: u16 = 1455;
741const CODEX_CALLBACK_FALLBACK_PORT: u16 = 1457;
742const CODEX_CALLBACK_PATH: &str = "/auth/callback";
743
744async fn do_oauth_flow(
745 http_client: Arc<dyn HttpClient>,
746 cx: &AsyncApp,
747) -> Result<CodexCredentials> {
748 // Start the callback server FIRST so the redirect URI is ready
749 let (redirect_uri, callback_rx) =
750 oauth_callback_server::start_oauth_callback_server_with_config(
751 oauth_callback_server::OAuthCallbackServerConfig {
752 host: CODEX_CALLBACK_HOST,
753 preferred_port: CODEX_CALLBACK_PORT,
754 fallback_port: Some(CODEX_CALLBACK_FALLBACK_PORT),
755 path: CODEX_CALLBACK_PATH,
756 },
757 )
758 .context("Failed to start OAuth callback server")?;
759
760 // PKCE verifier: 32 random bytes → base64url (no padding)
761 let mut verifier_bytes = [0u8; 32];
762 rand::rng().fill_bytes(&mut verifier_bytes);
763 let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
764
765 // PKCE challenge: SHA-256(verifier) → base64url
766 let mut hasher = Sha256::new();
767 hasher.update(verifier.as_bytes());
768 let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize().as_slice());
769
770 // CSRF state: 16 random bytes → hex string
771 let mut state_bytes = [0u8; 16];
772 rand::rng().fill_bytes(&mut state_bytes);
773 let oauth_state: String = state_bytes.iter().map(|b| format!("{b:02x}")).collect();
774
775 let mut auth_url = url::Url::parse(OPENAI_AUTHORIZE_URL).expect("valid base URL");
776 auth_url
777 .query_pairs_mut()
778 .append_pair("client_id", CLIENT_ID)
779 .append_pair("redirect_uri", &redirect_uri)
780 // Deliberately excludes `api.connectors.read api.connectors.invoke`
781 // (which Codex CLI requests): extra scopes inflate the
782 // access-token JWT, and the serialized credentials must fit within
783 // Windows Credential Manager's 2560-byte blob limit
784 // (CRED_MAX_CREDENTIAL_BLOB_SIZE). See #58541.
785 .append_pair("scope", "openid profile email offline_access")
786 .append_pair("response_type", "code")
787 .append_pair("code_challenge", &challenge)
788 .append_pair("code_challenge_method", "S256")
789 .append_pair("id_token_add_organizations", "true")
790 .append_pair("state", &oauth_state)
791 .append_pair("codex_cli_simplified_flow", "true")
792 .append_pair("originator", "zed");
793
794 // Open browser AFTER the listener is ready
795 cx.update(|cx| cx.open_url(auth_url.as_str()));
796
797 // Await the callback
798 let callback = callback_rx
799 .await
800 .map_err(|_| anyhow!("OAuth callback was cancelled"))?
801 .context("OAuth callback failed")?;
802
803 // Validate CSRF state
804 if callback.state != oauth_state {
805 return Err(anyhow!("OAuth state mismatch"));
806 }
807
808 let tokens = exchange_code(&http_client, &callback.code, &verifier, &redirect_uri)
809 .await
810 .context("Token exchange failed")?;
811
812 let jwt = tokens
813 .id_token
814 .as_deref()
815 .unwrap_or(tokens.access_token.as_str());
816 let claims = extract_jwt_claims(jwt);
817
818 Ok(CodexCredentials {
819 access_token: tokens.access_token,
820 refresh_token: tokens.refresh_token,
821 expires_at_ms: now_ms() + tokens.expires_in * 1000,
822 account_id: claims.account_id,
823 email: claims.email.or(tokens.email),
824 })
825}
826
827async fn exchange_code(
828 client: &Arc<dyn HttpClient>,
829 code: &str,
830 verifier: &str,
831 redirect_uri: &str,
832) -> Result<TokenResponse> {
833 let body = form_urlencoded::Serializer::new(String::new())
834 .append_pair("grant_type", "authorization_code")
835 .append_pair("client_id", CLIENT_ID)
836 .append_pair("code", code)
837 .append_pair("redirect_uri", redirect_uri)
838 .append_pair("code_verifier", verifier)
839 .finish();
840
841 let request = HttpRequest::builder()
842 .method(Method::POST)
843 .uri(OPENAI_TOKEN_URL)
844 .header("Content-Type", "application/x-www-form-urlencoded")
845 .body(AsyncBody::from(body))?;
846
847 let mut response = client.send(request).await?;
848 let mut body = String::new();
849 smol::io::AsyncReadExt::read_to_string(response.body_mut(), &mut body).await?;
850
851 if !response.status().is_success() {
852 return Err(anyhow!(
853 "Token exchange failed (HTTP {}): {body}",
854 response.status()
855 ));
856 }
857
858 serde_json::from_str::<TokenResponse>(&body).context("Failed to parse token response")
859}
860
861async fn refresh_token(
862 client: &Arc<dyn HttpClient>,
863 refresh_token: &str,
864) -> Result<CodexCredentials, RefreshError> {
865 let body = form_urlencoded::Serializer::new(String::new())
866 .append_pair("grant_type", "refresh_token")
867 .append_pair("client_id", CLIENT_ID)
868 .append_pair("refresh_token", refresh_token)
869 .finish();
870
871 let request = HttpRequest::builder()
872 .method(Method::POST)
873 .uri(OPENAI_TOKEN_URL)
874 .header("Content-Type", "application/x-www-form-urlencoded")
875 .body(AsyncBody::from(body))
876 .map_err(|e| RefreshError::Transient(e.into()))?;
877
878 let mut response = client
879 .send(request)
880 .await
881 .map_err(|e| RefreshError::Transient(e))?;
882 let status = response.status();
883 let mut body = String::new();
884 smol::io::AsyncReadExt::read_to_string(response.body_mut(), &mut body)
885 .await
886 .map_err(|e| RefreshError::Transient(e.into()))?;
887
888 if !status.is_success() {
889 let err = anyhow!("Token refresh failed (HTTP {}): {body}", status);
890 // 400/401/403 indicate a revoked or invalid refresh token.
891 // 5xx and other errors are treated as transient.
892 if status == http_client::StatusCode::BAD_REQUEST
893 || status == http_client::StatusCode::UNAUTHORIZED
894 || status == http_client::StatusCode::FORBIDDEN
895 {
896 return Err(RefreshError::Fatal(err));
897 }
898 return Err(RefreshError::Transient(err));
899 }
900
901 let tokens: TokenResponse =
902 serde_json::from_str(&body).map_err(|e| RefreshError::Transient(e.into()))?;
903 let jwt = tokens
904 .id_token
905 .as_deref()
906 .unwrap_or(tokens.access_token.as_str());
907 let claims = extract_jwt_claims(jwt);
908
909 Ok(CodexCredentials {
910 access_token: tokens.access_token,
911 refresh_token: tokens.refresh_token,
912 expires_at_ms: now_ms() + tokens.expires_in * 1000,
913 account_id: claims.account_id,
914 email: claims.email.or(tokens.email),
915 })
916}
917
918struct JwtClaims {
919 account_id: Option<String>,
920 email: Option<String>,
921}
922
923/// Extract claims from a JWT payload (base64url middle segment).
924/// Extracts `chatgpt_account_id` from three possible locations (matching Roo Code's
925/// implementation) and the `email` claim.
926fn extract_jwt_claims(jwt: &str) -> JwtClaims {
927 let Some(payload_b64) = jwt.split('.').nth(1) else {
928 return JwtClaims {
929 account_id: None,
930 email: None,
931 };
932 };
933 let Ok(payload) = URL_SAFE_NO_PAD.decode(payload_b64) else {
934 return JwtClaims {
935 account_id: None,
936 email: None,
937 };
938 };
939 let Ok(claims) = serde_json::from_slice::<serde_json::Value>(&payload) else {
940 return JwtClaims {
941 account_id: None,
942 email: None,
943 };
944 };
945
946 let account_id = claims
947 .get("chatgpt_account_id")
948 .and_then(|v| v.as_str())
949 .or_else(|| {
950 claims
951 .get("https://api.openai.com/auth")
952 .and_then(|v| v.get("chatgpt_account_id"))
953 .and_then(|v| v.as_str())
954 })
955 .or_else(|| {
956 claims
957 .get("organizations")
958 .and_then(|v| v.as_array())
959 .and_then(|arr| arr.first())
960 .and_then(|org| org.get("id"))
961 .and_then(|v| v.as_str())
962 })
963 .map(|s| s.to_owned());
964
965 let email = claims
966 .get("email")
967 .and_then(|v| v.as_str())
968 .map(|s| s.to_owned());
969
970 JwtClaims { account_id, email }
971}
972
973fn now_ms() -> u64 {
974 SystemTime::now()
975 .duration_since(UNIX_EPOCH)
976 .map(|d| d.as_millis() as u64)
977 .unwrap_or_else(|err| {
978 log::error!("System clock is before UNIX epoch: {err}");
979 0
980 })
981}
982
983fn do_sign_in(state: &Entity<State>, http_client: &Arc<dyn HttpClient>, cx: &mut App) {
984 if state.read(cx).is_signing_in() {
985 return;
986 }
987
988 let weak_state = state.downgrade();
989 let http_client = http_client.clone();
990
991 let task = cx.spawn(async move |cx| {
992 match do_oauth_flow(http_client, &*cx).await {
993 Ok(creds) => {
994 let persist_result = async {
995 let credentials_provider =
996 weak_state.read_with(&*cx, |s, _| s.credentials_provider.clone())?;
997 let json = serde_json::to_vec(&creds)?;
998 credentials_provider
999 .write_credentials(CREDENTIALS_KEY, "Bearer", &json, &*cx)
1000 .await?;
1001 anyhow::Ok(())
1002 }
1003 .await;
1004
1005 match persist_result {
1006 Ok(()) => {
1007 weak_state
1008 .update(cx, |s, cx| {
1009 s.credentials = Some(creds);
1010 s.sign_in_task = None;
1011 s.last_auth_error = None;
1012 cx.notify();
1013 })
1014 .log_err();
1015 }
1016 Err(err) => {
1017 log::error!(
1018 "ChatGPT subscription sign-in failed to persist credentials: {err:?}"
1019 );
1020 weak_state
1021 .update(cx, |s, cx| {
1022 s.sign_in_task = None;
1023 s.last_auth_error =
1024 Some("Failed to save credentials. Please try again.".into());
1025 cx.notify();
1026 })
1027 .log_err();
1028 }
1029 }
1030 }
1031 Err(err) => {
1032 log::error!("ChatGPT subscription sign-in failed: {err:?}");
1033 weak_state
1034 .update(cx, |s, cx| {
1035 s.sign_in_task = None;
1036 s.last_auth_error = Some("Sign-in failed. Please try again.".into());
1037 cx.notify();
1038 })
1039 .log_err();
1040 }
1041 }
1042 anyhow::Ok(())
1043 });
1044
1045 state.update(cx, |s, cx| {
1046 s.last_auth_error = None;
1047 s.sign_in_task = Some(task);
1048 cx.notify();
1049 });
1050}
1051
1052fn do_sign_out(state: &gpui::WeakEntity<State>, cx: &mut App) -> Task<Result<()>> {
1053 let weak_state = state.clone();
1054 // Clear credentials and cancel in-flight work immediately so the UI
1055 // reflects the sign-out right away.
1056 weak_state
1057 .update(cx, |s, cx| {
1058 s.auth_generation += 1;
1059 s.credentials = None;
1060 s.sign_in_task = None;
1061 s.refresh_task = None;
1062 s.last_auth_error = None;
1063 cx.notify();
1064 })
1065 .log_err();
1066
1067 cx.spawn(async move |cx| {
1068 let credentials_provider =
1069 weak_state.read_with(&*cx, |s, _| s.credentials_provider.clone())?;
1070 credentials_provider
1071 .delete_credentials(CREDENTIALS_KEY, &*cx)
1072 .await
1073 .context("Failed to delete ChatGPT subscription credentials from keychain")?;
1074 anyhow::Ok(())
1075 })
1076}
1077
1078struct ConfigurationView {
1079 state: Entity<State>,
1080 http_client: Arc<dyn HttpClient>,
1081 /// When `true`, the description is rendered elsewhere (the settings row's
1082 /// left column), so it's omitted here to avoid duplication.
1083 compact: bool,
1084}
1085
1086impl Render for ConfigurationView {
1087 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1088 let state = self.state.read(cx);
1089
1090 if state.is_authenticated() {
1091 let label = state
1092 .email()
1093 .map(|e| format!("Signed in as {e}"))
1094 .unwrap_or_else(|| "Signed in".to_string());
1095
1096 let weak_state = self.state.downgrade();
1097
1098 return v_flex()
1099 .child(
1100 ConfiguredApiCard::new("openai-subscribed-sign-out", SharedString::from(label))
1101 .button_label("Sign Out")
1102 .on_click(cx.listener(move |_this, _, _window, cx| {
1103 do_sign_out(&weak_state, cx).detach_and_log_err(cx);
1104 })),
1105 )
1106 .into_any_element();
1107 }
1108
1109 let last_auth_error = state.last_auth_error.clone();
1110 let provider_state = self.state.clone();
1111 let http_client = self.http_client.clone();
1112
1113 let is_signing_in = state.is_signing_in();
1114 let button_label = if is_signing_in {
1115 "Signing in…"
1116 } else {
1117 "Sign In"
1118 };
1119
1120 v_flex()
1121 .gap_2()
1122 .when(!self.compact, |this| {
1123 this.child(Label::new(SUBSCRIPTION_DESCRIPTION))
1124 })
1125 .child(
1126 Button::new("sign-in", button_label)
1127 .when(!self.compact, |this| this.full_width())
1128 .style(ButtonStyle::Outlined)
1129 .size(ButtonSize::Medium)
1130 .loading(is_signing_in)
1131 .disabled(is_signing_in)
1132 .on_click(move |_, _window, cx| {
1133 do_sign_in(&provider_state, &http_client, cx);
1134 }),
1135 )
1136 .when_some(last_auth_error, |this, error| {
1137 this.child(
1138 h_flex()
1139 .gap_1()
1140 .justify_center()
1141 .child(
1142 Icon::new(IconName::XCircle)
1143 .color(Color::Error)
1144 .size(IconSize::Small),
1145 )
1146 .child(Label::new(error).color(Color::Muted)),
1147 )
1148 })
1149 .into_any_element()
1150 }
1151}
1152
1153#[cfg(test)]
1154mod tests {
1155 use super::*;
1156 use gpui::TestAppContext;
1157 use http_client::FakeHttpClient;
1158 use parking_lot::Mutex;
1159 use std::future::Future;
1160 use std::pin::Pin;
1161 use std::sync::atomic::{AtomicUsize, Ordering};
1162
1163 struct FakeCredentialsProvider {
1164 storage: Mutex<Option<(String, Vec<u8>)>>,
1165 }
1166
1167 impl FakeCredentialsProvider {
1168 fn new() -> Self {
1169 Self {
1170 storage: Mutex::new(None),
1171 }
1172 }
1173 }
1174
1175 impl CredentialsProvider for FakeCredentialsProvider {
1176 fn read_credentials<'a>(
1177 &'a self,
1178 _url: &'a str,
1179 _cx: &'a AsyncApp,
1180 ) -> Pin<Box<dyn Future<Output = Result<Option<(String, Vec<u8>)>>> + 'a>> {
1181 Box::pin(async { Ok(self.storage.lock().clone()) })
1182 }
1183
1184 fn write_credentials<'a>(
1185 &'a self,
1186 _url: &'a str,
1187 username: &'a str,
1188 password: &'a [u8],
1189 _cx: &'a AsyncApp,
1190 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1191 self.storage
1192 .lock()
1193 .replace((username.to_string(), password.to_vec()));
1194 Box::pin(async { Ok(()) })
1195 }
1196
1197 fn delete_credentials<'a>(
1198 &'a self,
1199 _url: &'a str,
1200 _cx: &'a AsyncApp,
1201 ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
1202 *self.storage.lock() = None;
1203 Box::pin(async { Ok(()) })
1204 }
1205 }
1206
1207 fn make_expired_credentials() -> CodexCredentials {
1208 CodexCredentials {
1209 access_token: "old_access".to_string(),
1210 refresh_token: "old_refresh".to_string(),
1211 expires_at_ms: 0,
1212 account_id: None,
1213 email: None,
1214 }
1215 }
1216
1217 fn make_fresh_credentials() -> CodexCredentials {
1218 CodexCredentials {
1219 access_token: "fresh_access".to_string(),
1220 refresh_token: "fresh_refresh".to_string(),
1221 expires_at_ms: now_ms() + 3_600_000,
1222 account_id: None,
1223 email: None,
1224 }
1225 }
1226
1227 fn fake_token_response() -> String {
1228 serde_json::json!({
1229 "access_token": "fresh_access",
1230 "refresh_token": "fresh_refresh",
1231 "expires_in": 3600
1232 })
1233 .to_string()
1234 }
1235
1236 #[gpui::test]
1237 async fn test_concurrent_refresh_deduplicates(cx: &mut TestAppContext) {
1238 let refresh_count = Arc::new(AtomicUsize::new(0));
1239 let refresh_count_clone = refresh_count.clone();
1240
1241 let http_client = FakeHttpClient::create(move |_request| {
1242 let refresh_count = refresh_count_clone.clone();
1243 async move {
1244 refresh_count.fetch_add(1, Ordering::SeqCst);
1245 let body = fake_token_response();
1246 Ok(http_client::Response::builder()
1247 .status(200)
1248 .body(http_client::AsyncBody::from(body))?)
1249 }
1250 });
1251
1252 let state = cx.new(|_cx| State {
1253 credentials: Some(make_expired_credentials()),
1254 sign_in_task: None,
1255 refresh_task: None,
1256 load_task: None,
1257 credentials_provider: Arc::new(FakeCredentialsProvider::new()),
1258 auth_generation: 0,
1259 last_auth_error: None,
1260 });
1261
1262 let weak_state = cx.read(|_cx| state.downgrade());
1263 let http: Arc<dyn HttpClient> = http_client;
1264
1265 // Spawn two concurrent refresh attempts.
1266 let weak1 = weak_state.clone();
1267 let http1 = http.clone();
1268 let task1 =
1269 cx.spawn(async move |mut cx| get_fresh_credentials(&weak1, &http1, &mut cx).await);
1270
1271 let weak2 = weak_state.clone();
1272 let http2 = http.clone();
1273 let task2 =
1274 cx.spawn(async move |mut cx| get_fresh_credentials(&weak2, &http2, &mut cx).await);
1275
1276 // Drive both to completion.
1277 cx.run_until_parked();
1278 let result1 = task1.await;
1279 let result2 = task2.await;
1280
1281 assert!(result1.is_ok(), "first refresh should succeed");
1282 assert!(result2.is_ok(), "second refresh should succeed");
1283 assert_eq!(result1.unwrap().access_token, "fresh_access");
1284 assert_eq!(result2.unwrap().access_token, "fresh_access");
1285 assert_eq!(
1286 refresh_count.load(Ordering::SeqCst),
1287 1,
1288 "refresh_token should only be called once despite two concurrent callers"
1289 );
1290 }
1291
1292 #[gpui::test]
1293 async fn test_fresh_credentials_skip_refresh(cx: &mut TestAppContext) {
1294 let refresh_count = Arc::new(AtomicUsize::new(0));
1295 let refresh_count_clone = refresh_count.clone();
1296
1297 let http_client = FakeHttpClient::create(move |_request| {
1298 let refresh_count = refresh_count_clone.clone();
1299 async move {
1300 refresh_count.fetch_add(1, Ordering::SeqCst);
1301 let body = fake_token_response();
1302 Ok(http_client::Response::builder()
1303 .status(200)
1304 .body(http_client::AsyncBody::from(body))?)
1305 }
1306 });
1307
1308 let state = cx.new(|_cx| State {
1309 credentials: Some(make_fresh_credentials()),
1310 sign_in_task: None,
1311 refresh_task: None,
1312 load_task: None,
1313 credentials_provider: Arc::new(FakeCredentialsProvider::new()),
1314 auth_generation: 0,
1315 last_auth_error: None,
1316 });
1317
1318 let weak_state = cx.read(|_cx| state.downgrade());
1319 let http: Arc<dyn HttpClient> = http_client;
1320
1321 let weak = weak_state.clone();
1322 let http_clone = http.clone();
1323 let result = cx
1324 .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await)
1325 .await;
1326
1327 assert!(result.is_ok());
1328 assert_eq!(result.unwrap().access_token, "fresh_access");
1329 assert_eq!(
1330 refresh_count.load(Ordering::SeqCst),
1331 0,
1332 "no refresh should happen when credentials are fresh"
1333 );
1334 }
1335
1336 #[gpui::test]
1337 async fn test_no_credentials_returns_no_api_key(cx: &mut TestAppContext) {
1338 let http_client = FakeHttpClient::create(|_| async {
1339 Ok(http_client::Response::builder()
1340 .status(200)
1341 .body(http_client::AsyncBody::default())?)
1342 });
1343
1344 let state = cx.new(|_cx| State {
1345 credentials: None,
1346 sign_in_task: None,
1347 refresh_task: None,
1348 load_task: None,
1349 credentials_provider: Arc::new(FakeCredentialsProvider::new()),
1350 auth_generation: 0,
1351 last_auth_error: None,
1352 });
1353
1354 let weak_state = cx.read(|_cx| state.downgrade());
1355 let http: Arc<dyn HttpClient> = http_client;
1356
1357 let weak = weak_state.clone();
1358 let http_clone = http.clone();
1359 let result = cx
1360 .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await)
1361 .await;
1362
1363 assert!(matches!(
1364 result,
1365 Err(LanguageModelCompletionError::NoApiKey { .. })
1366 ));
1367 }
1368
1369 #[gpui::test]
1370 async fn test_fatal_refresh_clears_auth_state(cx: &mut TestAppContext) {
1371 let http_client = FakeHttpClient::create(move |_request| async move {
1372 Ok(http_client::Response::builder()
1373 .status(401)
1374 .body(http_client::AsyncBody::from(r#"{"error":"invalid_grant"}"#))?)
1375 });
1376
1377 let creds_provider = Arc::new(FakeCredentialsProvider::new());
1378 let state = cx.new(|_cx| State {
1379 credentials: Some(make_expired_credentials()),
1380 sign_in_task: None,
1381 refresh_task: None,
1382 load_task: None,
1383 credentials_provider: creds_provider.clone(),
1384 auth_generation: 0,
1385 last_auth_error: None,
1386 });
1387
1388 let weak_state = cx.read(|_cx| state.downgrade());
1389 let http: Arc<dyn HttpClient> = http_client;
1390
1391 let weak = weak_state.clone();
1392 let http_clone = http.clone();
1393 let result = cx
1394 .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await)
1395 .await;
1396
1397 cx.run_until_parked();
1398
1399 assert!(result.is_err(), "fatal refresh should return an error");
1400 cx.read(|cx| {
1401 let s = state.read(cx);
1402 assert!(
1403 s.credentials.is_none(),
1404 "credentials should be cleared on fatal refresh failure"
1405 );
1406 assert!(
1407 s.last_auth_error.is_some(),
1408 "last_auth_error should be set on fatal refresh failure"
1409 );
1410 });
1411 }
1412
1413 #[gpui::test]
1414 async fn test_transient_refresh_keeps_credentials(cx: &mut TestAppContext) {
1415 let http_client = FakeHttpClient::create(move |_request| async move {
1416 Ok(http_client::Response::builder()
1417 .status(500)
1418 .body(http_client::AsyncBody::from("Internal Server Error"))?)
1419 });
1420
1421 let state = cx.new(|_cx| State {
1422 credentials: Some(make_expired_credentials()),
1423 sign_in_task: None,
1424 refresh_task: None,
1425 load_task: None,
1426 credentials_provider: Arc::new(FakeCredentialsProvider::new()),
1427 auth_generation: 0,
1428 last_auth_error: None,
1429 });
1430
1431 let weak_state = cx.read(|_cx| state.downgrade());
1432 let http: Arc<dyn HttpClient> = http_client;
1433
1434 let weak = weak_state.clone();
1435 let http_clone = http.clone();
1436 let result = cx
1437 .spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await)
1438 .await;
1439
1440 cx.run_until_parked();
1441
1442 assert!(result.is_err(), "transient refresh should return an error");
1443 cx.read(|cx| {
1444 let s = state.read(cx);
1445 assert!(
1446 s.credentials.is_some(),
1447 "credentials should be kept on transient refresh failure"
1448 );
1449 assert!(
1450 s.last_auth_error.is_none(),
1451 "last_auth_error should not be set on transient refresh failure"
1452 );
1453 });
1454 }
1455
1456 #[gpui::test]
1457 async fn test_sign_out_during_refresh_discards_result(cx: &mut TestAppContext) {
1458 let (gate_tx, gate_rx) = futures::channel::oneshot::channel::<()>();
1459 let gate_rx = Arc::new(Mutex::new(Some(gate_rx)));
1460 let gate_rx_clone = gate_rx.clone();
1461
1462 let http_client = FakeHttpClient::create(move |_request| {
1463 let gate_rx = gate_rx_clone.clone();
1464 async move {
1465 // Wait until the gate is opened, simulating a slow network.
1466 let rx = gate_rx.lock().take();
1467 if let Some(rx) = rx {
1468 let _ = rx.await;
1469 }
1470 let body = fake_token_response();
1471 Ok(http_client::Response::builder()
1472 .status(200)
1473 .body(http_client::AsyncBody::from(body))?)
1474 }
1475 });
1476
1477 let creds_provider = Arc::new(FakeCredentialsProvider::new());
1478 let state = cx.new(|_cx| State {
1479 credentials: Some(make_expired_credentials()),
1480 sign_in_task: None,
1481 refresh_task: None,
1482 load_task: None,
1483 credentials_provider: creds_provider.clone(),
1484 auth_generation: 0,
1485 last_auth_error: None,
1486 });
1487
1488 let weak_state = cx.read(|_cx| state.downgrade());
1489 let http: Arc<dyn HttpClient> = http_client;
1490
1491 // Start a refresh
1492 let weak = weak_state.clone();
1493 let http_clone = http.clone();
1494 let refresh_task =
1495 cx.spawn(async move |mut cx| get_fresh_credentials(&weak, &http_clone, &mut cx).await);
1496
1497 cx.run_until_parked();
1498
1499 // Sign out while the refresh is in-flight
1500 cx.update(|cx| {
1501 do_sign_out(&weak_state, cx).detach();
1502 });
1503 cx.run_until_parked();
1504
1505 // Now let the refresh respond by opening the gate
1506 let _ = gate_tx.send(());
1507 cx.run_until_parked();
1508
1509 let result = refresh_task.await;
1510 assert!(result.is_err(), "refresh should fail after sign-out");
1511
1512 cx.read(|cx| {
1513 let s = state.read(cx);
1514 assert!(
1515 s.credentials.is_none(),
1516 "sign-out should have cleared credentials"
1517 );
1518 });
1519 }
1520
1521 #[gpui::test]
1522 async fn test_sign_out_completes_fully(cx: &mut TestAppContext) {
1523 let creds_provider = Arc::new(FakeCredentialsProvider::new());
1524 // Pre-populate the credential store
1525 creds_provider
1526 .storage
1527 .lock()
1528 .replace(("Bearer".to_string(), b"some-creds".to_vec()));
1529
1530 let state = cx.new(|_cx| State {
1531 credentials: Some(make_fresh_credentials()),
1532 sign_in_task: None,
1533 refresh_task: None,
1534 load_task: None,
1535 credentials_provider: creds_provider.clone(),
1536 auth_generation: 0,
1537 last_auth_error: None,
1538 });
1539
1540 let weak_state = cx.read(|_cx| state.downgrade());
1541 let sign_out_task = cx.update(|cx| do_sign_out(&weak_state, cx));
1542
1543 cx.run_until_parked();
1544 sign_out_task.await.expect("sign-out should succeed");
1545
1546 assert!(
1547 creds_provider.storage.lock().is_none(),
1548 "credential store should be empty after sign-out"
1549 );
1550 cx.read(|cx| {
1551 assert!(
1552 !state.read(cx).is_authenticated(),
1553 "state should show not authenticated"
1554 );
1555 });
1556 }
1557
1558 #[gpui::test]
1559 async fn test_authenticate_awaits_initial_load(cx: &mut TestAppContext) {
1560 let creds = make_fresh_credentials();
1561 let creds_json = serde_json::to_vec(&creds).unwrap();
1562 let creds_provider = Arc::new(FakeCredentialsProvider::new());
1563 creds_provider
1564 .storage
1565 .lock()
1566 .replace(("Bearer".to_string(), creds_json));
1567
1568 let http_client = FakeHttpClient::create(|_| async {
1569 Ok(http_client::Response::builder()
1570 .status(200)
1571 .body(http_client::AsyncBody::default())?)
1572 });
1573
1574 let provider =
1575 cx.update(|cx| OpenAiSubscribedProvider::new(http_client, creds_provider, cx));
1576
1577 // Before load completes, authenticate should still await the load.
1578 let auth_task = cx.update(|cx| provider.authenticate(cx));
1579
1580 // Drive the load to completion.
1581 cx.run_until_parked();
1582
1583 let result = auth_task.await;
1584 assert!(
1585 result.is_ok(),
1586 "authenticate should succeed after load completes with valid credentials"
1587 );
1588 }
1589}
1590