valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add AuthManager and enhance GetAuthStatus command (#2577)

Browse Source 
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:
Eric Traut
2025-08-22 13:10:11 -07:00
committed by GitHub
parent cdc77c10fb
commit dc42ec0eb4
26 changed files with 674 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
codex-rs/core/src/client.rs
View File

@@ -4,8 +4,8 @@ use std::sync::OnceLock;
use std::time::Duration;
use bytes::Bytes;
use codex_login::AuthManager;
use codex_login::AuthMode;
use codex_login::CodexAuth;
use eventsource_stream::Eventsource;
use futures::prelude::*;
use regex_lite::Regex;
@@ -61,7 +61,7 @@ struct Error {
#[derive(Debug, Clone)]
pub struct ModelClient {
config: Arc<Config>,
auth: Option<CodexAuth>,
auth_manager: Option<Arc<AuthManager>>,
client: reqwest::Client,
provider: ModelProviderInfo,
session_id: Uuid,
@@ -72,7 +72,7 @@ pub struct ModelClient {
impl ModelClient {
pub fn new(
config: Arc<Config>,
auth: Option<CodexAuth>,
auth_manager: Option<Arc<AuthManager>>,
provider: ModelProviderInfo,
effort: ReasoningEffortConfig,
summary: ReasoningSummaryConfig,
@@ -80,7 +80,7 @@ impl ModelClient {
) -> Self {
Self {
config,
auth,
auth_manager,
client: reqwest::Client::new(),
provider,
session_id,
@@ -141,7 +141,8 @@ impl ModelClient {
return stream_from_fixture(path, self.provider.clone()).await;
}
let auth = self.auth.clone();
let auth_manager = self.auth_manager.clone();
let auth = auth_manager.as_ref().and_then(|m| m.auth());
let auth_mode = auth.as_ref().map(|a| a.mode);
@@ -264,9 +265,10 @@ impl ModelClient {
.and_then(|s| s.parse::<u64>().ok());
if status == StatusCode::UNAUTHORIZED
&& let Some(a) = auth.as_ref()
&& let Some(manager) = auth_manager.as_ref()
&& manager.auth().is_some()
{
let _ = a.refresh_token().await;
let _ = manager.refresh_token().await;
}
// The OpenAI Responses endpoint returns structured JSON bodies even for 4xx/5xx
@@ -353,8 +355,8 @@ impl ModelClient {
self.summary
}
pub fn get_auth(&self) -> Option<CodexAuth> {
self.auth.clone()
pub fn get_auth_manager(&self) -> Option<Arc<AuthManager>> {
self.auth_manager.clone()
}
}

34
codex-rs/core/src/codex.rs
View File

@@ -13,7 +13,7 @@ use async_channel::Sender;
use codex_apply_patch::ApplyPatchAction;
use codex_apply_patch::MaybeApplyPatchVerified;
use codex_apply_patch::maybe_parse_apply_patch_verified;
use codex_login::CodexAuth;
use codex_login::AuthManager;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use futures::prelude::*;
@@ -144,7 +144,10 @@ pub(crate) const INITIAL_SUBMIT_ID: &str = "";
impl Codex {
/// Spawn a new [`Codex`] and initialize the session.
pub async fn spawn(config: Config, auth: Option<CodexAuth>) -> CodexResult<CodexSpawnOk> {
pub async fn spawn(
config: Config,
auth_manager: Arc<AuthManager>,
) -> CodexResult<CodexSpawnOk> {
let (tx_sub, rx_sub) = async_channel::bounded(64);
let (tx_event, rx_event) = async_channel::unbounded();
@@ -169,13 +172,17 @@ impl Codex {
};
// Generate a unique ID for the lifetime of this Codex session.
let (session, turn_context) =
Session::new(configure_session, config.clone(), auth, tx_event.clone())
.await
.map_err(|e| {
error!("Failed to create session: {e:#}");
CodexErr::InternalAgentDied
})?;
let (session, turn_context) = Session::new(
configure_session,
config.clone(),
auth_manager.clone(),
tx_event.clone(),
)
.await
.map_err(|e| {
error!("Failed to create session: {e:#}");
CodexErr::InternalAgentDied
})?;
let session_id = session.session_id;
// This task will run until Op::Shutdown is received.
@@ -323,7 +330,7 @@ impl Session {
async fn new(
configure_session: ConfigureSession,
config: Arc<Config>,
auth: Option<CodexAuth>,
auth_manager: Arc<AuthManager>,
tx_event: Sender<Event>,
) -> anyhow::Result<(Arc<Self>, TurnContext)> {
let ConfigureSession {
@@ -467,7 +474,7 @@ impl Session {
// construct the model client.
let client = ModelClient::new(
config.clone(),
auth.clone(),
Some(auth_manager.clone()),
provider.clone(),
model_reasoning_effort,
model_reasoning_summary,
@@ -1034,7 +1041,8 @@ async fn submission_loop(
let effective_effort = effort.unwrap_or(prev.client.get_reasoning_effort());
let effective_summary = summary.unwrap_or(prev.client.get_reasoning_summary());
let auth = prev.client.get_auth();
let auth_manager = prev.client.get_auth_manager();
// Build updated config for the client
let mut updated_config = (*config).clone();
updated_config.model = effective_model.clone();
@@ -1042,7 +1050,7 @@ async fn submission_loop(
let client = ModelClient::new(
Arc::new(updated_config),
auth,
auth_manager,
provider,
effective_effort,
effective_summary,

37
codex-rs/core/src/conversation_manager.rs
View File

@@ -1,6 +1,7 @@
use std::collections::HashMap;
use std::sync::Arc;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use tokio::sync::RwLock;
use uuid::Uuid;
@@ -28,33 +29,37 @@ pub struct NewConversation {
/// maintaining them in memory.
pub struct ConversationManager {
conversations: Arc<RwLock<HashMap<Uuid, Arc<CodexConversation>>>>,
}
impl Default for ConversationManager {
fn default() -> Self {
Self {
conversations: Arc::new(RwLock::new(HashMap::new())),
}
}
auth_manager: Arc<AuthManager>,
}
impl ConversationManager {
pub async fn new_conversation(&self, config: Config) -> CodexResult<NewConversation> {
let auth = CodexAuth::from_codex_home(&config.codex_home, config.preferred_auth_method)?;
self.new_conversation_with_auth(config, auth).await
pub fn new(auth_manager: Arc<AuthManager>) -> Self {
Self {
conversations: Arc::new(RwLock::new(HashMap::new())),
auth_manager,
}
}
/// Used for integration tests: should not be used by ordinary business
/// logic.
pub async fn new_conversation_with_auth(
/// Construct with a dummy AuthManager containing the provided CodexAuth.
/// Used for integration tests: should not be used by ordinary business logic.
pub fn with_auth(auth: CodexAuth) -> Self {
Self::new(codex_login::AuthManager::from_auth_for_testing(auth))
}
pub async fn new_conversation(&self, config: Config) -> CodexResult<NewConversation> {
self.spawn_conversation(config, self.auth_manager.clone())
.await
}
async fn spawn_conversation(
&self,
config: Config,
auth: Option<CodexAuth>,
auth_manager: Arc<AuthManager>,
) -> CodexResult<NewConversation> {
let CodexSpawnOk {
codex,
session_id: conversation_id,
} = Codex::spawn(config, auth).await?;
} = Codex::spawn(config, auth_manager).await?;
// The first event must be `SessionInitialized`. Validate and forward it
// to the caller so that they can display it in the conversation