Add AuthManager and enhance GetAuthStatus command (#2577)
This PR adds a central `AuthManager` struct that manages the auth information used across conversations and the MCP server. Prior to this, each conversation and the MCP server got their own private snapshots of the auth information, and changes to one (such as a logout or token refresh) were not seen by others. This is especially problematic when multiple instances of the CLI are run. For example, consider the case where you start CLI 1 and log in to ChatGPT account X and then start CLI 2 and log out and then log in to ChatGPT account Y. The conversation in CLI 1 is still using account X, but if you create a new conversation, it will suddenly (and unexpectedly) switch to account Y. With the `AuthManager`, auth information is read from disk at the time the `ConversationManager` is constructed, and it is cached in memory. All new conversations use this same auth information, as do any token refreshes. The `AuthManager` is also used by the MCP server's GetAuthStatus command, which now returns the auth method currently used by the MCP server. This PR also includes an enhancement to the GetAuthStatus command. It now accepts two new (optional) input parameters: `include_token` and `refresh_token`. Callers can use this to request the in-use auth token and can optionally request to refresh the token. The PR also adds tests for the login and auth APIs that I recently added to the MCP server.
This commit is contained in:
@@ -9,6 +9,7 @@ use codex_ansi_escape::ansi_escape_line;
|
||||
use codex_core::ConversationManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::protocol::TokenUsage;
|
||||
use codex_login::AuthManager;
|
||||
use color_eyre::eyre::Result;
|
||||
use crossterm::event::KeyCode;
|
||||
use crossterm::event::KeyEvent;
|
||||
@@ -50,6 +51,7 @@ pub(crate) struct App {
|
||||
impl App {
|
||||
pub async fn run(
|
||||
tui: &mut tui::Tui,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
config: Config,
|
||||
initial_prompt: Option<String>,
|
||||
initial_images: Vec<PathBuf>,
|
||||
@@ -58,7 +60,7 @@ impl App {
|
||||
let (app_event_tx, mut app_event_rx) = unbounded_channel();
|
||||
let app_event_tx = AppEventSender::new(app_event_tx);
|
||||
|
||||
let conversation_manager = Arc::new(ConversationManager::default());
|
||||
let conversation_manager = Arc::new(ConversationManager::new(auth_manager.clone()));
|
||||
|
||||
let enhanced_keys_supported = supports_keyboard_enhancement().unwrap_or(false);
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@ use codex_core::protocol::PatchApplyBeginEvent;
|
||||
use codex_core::protocol::PatchApplyEndEvent;
|
||||
use codex_core::protocol::StreamErrorEvent;
|
||||
use codex_core::protocol::TaskCompleteEvent;
|
||||
use codex_login::CodexAuth;
|
||||
use crossterm::event::KeyCode;
|
||||
use crossterm::event::KeyEvent;
|
||||
use crossterm::event::KeyModifiers;
|
||||
@@ -104,7 +105,9 @@ async fn helpers_are_available_and_do_not_panic() {
|
||||
let (tx_raw, _rx) = unbounded_channel::<AppEvent>();
|
||||
let tx = AppEventSender::new(tx_raw);
|
||||
let cfg = test_config();
|
||||
let conversation_manager = Arc::new(ConversationManager::default());
|
||||
let conversation_manager = Arc::new(ConversationManager::with_auth(CodexAuth::from_api_key(
|
||||
"test",
|
||||
)));
|
||||
let mut w = ChatWidget::new(
|
||||
cfg,
|
||||
conversation_manager,
|
||||
|
||||
@@ -12,6 +12,7 @@ use codex_core::config::find_codex_home;
|
||||
use codex_core::config::load_config_as_toml_with_cli_overrides;
|
||||
use codex_core::protocol::AskForApproval;
|
||||
use codex_core::protocol::SandboxPolicy;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::AuthMode;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_ollama::DEFAULT_OSS_MODEL;
|
||||
@@ -300,6 +301,7 @@ async fn run_ratatui_app(
|
||||
|
||||
let Cli { prompt, images, .. } = cli;
|
||||
|
||||
let auth_manager = AuthManager::shared(config.codex_home.clone(), config.preferred_auth_method);
|
||||
let login_status = get_login_status(&config);
|
||||
let should_show_onboarding =
|
||||
should_show_onboarding(login_status, &config, should_show_trust_screen);
|
||||
@@ -312,6 +314,7 @@ async fn run_ratatui_app(
|
||||
show_trust_screen: should_show_trust_screen,
|
||||
login_status,
|
||||
preferred_auth_method: config.preferred_auth_method,
|
||||
auth_manager: auth_manager.clone(),
|
||||
},
|
||||
&mut tui,
|
||||
)
|
||||
@@ -322,7 +325,7 @@ async fn run_ratatui_app(
|
||||
}
|
||||
}
|
||||
|
||||
let app_result = App::run(&mut tui, config, prompt, images).await;
|
||||
let app_result = App::run(&mut tui, auth_manager, config, prompt, images).await;
|
||||
|
||||
restore();
|
||||
// Mark the end of the recorded session.
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#![allow(clippy::unwrap_used)]
|
||||
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CLIENT_ID;
|
||||
use codex_login::ServerOptions;
|
||||
use codex_login::ShutdownHandle;
|
||||
@@ -112,6 +113,7 @@ pub(crate) struct AuthModeWidget {
|
||||
pub codex_home: PathBuf,
|
||||
pub login_status: LoginStatus,
|
||||
pub preferred_auth_method: AuthMode,
|
||||
pub auth_manager: Arc<AuthManager>,
|
||||
}
|
||||
|
||||
impl AuthModeWidget {
|
||||
@@ -338,6 +340,7 @@ impl AuthModeWidget {
|
||||
Ok(child) => {
|
||||
let sign_in_state = self.sign_in_state.clone();
|
||||
let request_frame = self.request_frame.clone();
|
||||
let auth_manager = self.auth_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
let auth_url = child.auth_url.clone();
|
||||
{
|
||||
@@ -351,6 +354,9 @@ impl AuthModeWidget {
|
||||
let r = child.block_until_done().await;
|
||||
match r {
|
||||
Ok(()) => {
|
||||
// Force the auth manager to reload the new auth information.
|
||||
auth_manager.reload();
|
||||
|
||||
*sign_in_state.write().unwrap() = SignInState::ChatGptSuccessMessage;
|
||||
request_frame.schedule_frame();
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use codex_core::util::is_inside_git_repo;
|
||||
use codex_login::AuthManager;
|
||||
use crossterm::event::KeyCode;
|
||||
use crossterm::event::KeyEvent;
|
||||
use crossterm::event::KeyEventKind;
|
||||
@@ -58,6 +59,7 @@ pub(crate) struct OnboardingScreenArgs {
|
||||
pub show_login_screen: bool,
|
||||
pub login_status: LoginStatus,
|
||||
pub preferred_auth_method: AuthMode,
|
||||
pub auth_manager: Arc<AuthManager>,
|
||||
}
|
||||
|
||||
impl OnboardingScreen {
|
||||
@@ -69,6 +71,7 @@ impl OnboardingScreen {
|
||||
show_login_screen,
|
||||
login_status,
|
||||
preferred_auth_method,
|
||||
auth_manager,
|
||||
} = args;
|
||||
let mut steps: Vec<Step> = vec![Step::Welcome(WelcomeWidget {
|
||||
is_logged_in: !matches!(login_status, LoginStatus::NotAuthenticated),
|
||||
@@ -82,6 +85,7 @@ impl OnboardingScreen {
|
||||
codex_home: codex_home.clone(),
|
||||
login_status,
|
||||
preferred_auth_method,
|
||||
auth_manager,
|
||||
}))
|
||||
}
|
||||
let is_git_repo = is_inside_git_repo(&cwd);
|
||||
|
||||
Reference in New Issue
Block a user