From 1525bbdb9a77ce7a817b19465a6cce8e2c0800bf Mon Sep 17 00:00:00 2001 From: Ruslan Nigmatullin Date: Mon, 6 Apr 2026 12:46:55 -0700 Subject: [PATCH] app-server: centralize AuthManager initialization (#16764) Extract a shared helper that builds AuthManager from Config and applies the forced ChatGPT workspace override in one place. Create the shared AuthManager at MessageProcessor call sites so that upcoming new transport's initialization can reuse the same handle, and keep only external auth refresher wiring inside `MessageProcessor`. Remove the now-unused `AuthManager::shared_with_external_auth` helper. --- codex-rs/app-server/src/in_process.rs | 5 ++- codex-rs/app-server/src/lib.rs | 11 +++--- codex-rs/app-server/src/message_processor.rs | 16 +++------ .../src/message_processor/tracing_tests.rs | 5 ++- codex-rs/core/src/config/mod.rs | 15 ++++++++ codex-rs/core/src/connectors.rs | 17 ++++------ codex-rs/core/src/prompt_debug.rs | 8 ++--- codex-rs/login/src/auth/manager.rs | 34 ++++++++++++++----- codex-rs/login/src/lib.rs | 1 + codex-rs/mcp-server/src/message_processor.rs | 5 ++- 10 files changed, 69 insertions(+), 48 deletions(-) diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index eda2c75da..71beb58dc 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -80,6 +80,7 @@ use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::LoaderOverrides; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; +use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -379,6 +380,8 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { }); let processor_outgoing = Arc::clone(&outgoing_message_sender); + let auth_manager = + AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env); let (processor_tx, mut processor_rx) = mpsc::channel::(channel_capacity); let mut processor_handle = tokio::spawn(async move { let mut processor = MessageProcessor::new(MessageProcessorArgs { @@ -393,7 +396,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { log_db: None, config_warnings: args.config_warnings, session_source: args.session_source, - enable_codex_api_key_env: args.enable_codex_api_key_env, + auth_manager, rpc_transport: AppServerRpcTransport::InProcess, }); let mut thread_created_rx = processor.thread_created_receiver(); diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 8c3742adc..14250c205 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -396,11 +396,8 @@ pub async fn run_main_with_transport( } } - let auth_manager = AuthManager::shared( - config.codex_home.clone(), - /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, - ); + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); cloud_requirements_loader( auth_manager, config.chatgpt_base_url, @@ -611,6 +608,8 @@ pub async fn run_main_with_transport( let processor_handle = tokio::spawn({ let outgoing_message_sender = Arc::new(OutgoingMessageSender::new(outgoing_tx)); let outbound_control_tx = outbound_control_tx; + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); let cli_overrides: Vec<(String, TomlValue)> = cli_kv_overrides.clone(); let loader_overrides = loader_overrides_for_config_api; let mut processor = MessageProcessor::new(MessageProcessorArgs { @@ -625,7 +624,7 @@ pub async fn run_main_with_transport( log_db, config_warnings, session_source, - enable_codex_api_key_env: false, + auth_manager, rpc_transport: analytics_rpc_transport(transport), }); let mut thread_created_rx = processor.thread_created_receiver(); diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 02b948bc1..8e3c5be81 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -193,7 +193,7 @@ pub(crate) struct MessageProcessorArgs { pub(crate) log_db: Option, pub(crate) config_warnings: Vec, pub(crate) session_source: SessionSource, - pub(crate) enable_codex_api_key_env: bool, + pub(crate) auth_manager: Arc, pub(crate) rpc_transport: AppServerRpcTransport, } @@ -213,17 +213,12 @@ impl MessageProcessor { log_db, config_warnings, session_source, - enable_codex_api_key_env, + auth_manager, rpc_transport, } = args; - let auth_manager = AuthManager::shared_with_external_auth( - config.codex_home.clone(), - enable_codex_api_key_env, - config.cli_auth_credentials_store_mode, - Arc::new(ExternalAuthRefreshBridge { - outgoing: outgoing.clone(), - }), - ); + auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { + outgoing: outgoing.clone(), + })); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager.clone(), @@ -235,7 +230,6 @@ impl MessageProcessor { }, environment_manager, )); - auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id.clone()); let analytics_events_client = AnalyticsEventsClient::new( Arc::clone(&auth_manager), config.chatgpt_base_url.trim_end_matches('/').to_string(), diff --git a/codex-rs/app-server/src/message_processor/tracing_tests.rs b/codex-rs/app-server/src/message_processor/tracing_tests.rs index d110f4fee..2e8781606 100644 --- a/codex-rs/app-server/src/message_processor/tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor/tracing_tests.rs @@ -27,6 +27,7 @@ use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::LoaderOverrides; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; +use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use opentelemetry::global; @@ -234,6 +235,8 @@ fn build_test_processor( ) { let (outgoing_tx, outgoing_rx) = mpsc::channel(16); let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); + let auth_manager = + AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false); let processor = MessageProcessor::new(MessageProcessorArgs { outgoing, arg0_paths: Arg0DispatchPaths::default(), @@ -246,7 +249,7 @@ fn build_test_processor( log_db: None, config_warnings: Vec::new(), session_source: SessionSource::VSCode, - enable_codex_api_key_env: false, + auth_manager, rpc_transport: AppServerRpcTransport::Stdio, }); (processor, outgoing_rx) diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index d83fc8737..cb497815f 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -58,6 +58,7 @@ use codex_features::Features; use codex_features::FeaturesToml; use codex_git_utils::resolve_root_git_project_for_trust; use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthManagerConfig; use codex_mcp::mcp::McpConfig; use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::LMSTUDIO_OSS_PROVIDER_ID; @@ -593,6 +594,20 @@ pub struct Config { pub otel: codex_config::types::OtelConfig, } +impl AuthManagerConfig for Config { + fn codex_home(&self) -> PathBuf { + self.codex_home.clone() + } + + fn cli_auth_credentials_store_mode(&self) -> AuthCredentialsStoreMode { + self.cli_auth_credentials_store_mode + } + + fn forced_chatgpt_workspace_id(&self) -> Option { + self.forced_chatgpt_workspace_id.clone() + } +} + #[derive(Debug, Clone, Default)] pub struct ConfigBuilder { codex_home: Option, diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 10e10e99f..66ba6a406 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -143,7 +143,8 @@ pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth( pub async fn list_cached_accessible_connectors_from_mcp_tools( config: &Config, ) -> Option> { - let auth_manager = auth_manager_from_config(config); + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); let auth = auth_manager.auth().await; if !config.features.apps_enabled_for_auth(auth.as_ref()) { return Some(Vec::new()); @@ -182,7 +183,8 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( config: &Config, force_refetch: bool, ) -> anyhow::Result { - let auth_manager = auth_manager_from_config(config); + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); let auth = auth_manager.auth().await; if !config.features.apps_enabled_for_auth(auth.as_ref()) { return Ok(AccessibleConnectorsStatus { @@ -417,7 +419,8 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( let token_data = if let Some(auth) = auth { auth.get_token_data().ok() } else { - let auth_manager = auth_manager_from_config(config); + let auth_manager = + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); auth_manager .auth() .await @@ -492,14 +495,6 @@ async fn chatgpt_get_request_with_token( } } -fn auth_manager_from_config(config: &Config) -> std::sync::Arc { - AuthManager::shared( - config.codex_home.clone(), - /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, - ) -} - pub fn connector_display_label(connector: &AppInfo) -> String { format_connector_label(&connector.name, &connector.id) } diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 4a8513bb4..5fe39ee88 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -26,12 +26,8 @@ pub async fn build_prompt_input( ) -> CodexResult> { config.ephemeral = true; - let auth_manager = AuthManager::shared( - config.codex_home.clone(), - /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, - ); - auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id.clone()); + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); let thread_manager = ThreadManager::new( &config, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9f58f0f09..f0b31f0f4 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1108,6 +1108,23 @@ pub struct AuthManager { external_auth: RwLock>>, } +/// Configuration view required to construct a shared [`AuthManager`]. +/// +/// Implementations should return the auth-related config values for the +/// already-resolved runtime configuration. The primary implementation is +/// `codex_core::config::Config`, but this trait keeps `codex-login` independent +/// from `codex-core`. +pub trait AuthManagerConfig { + /// Returns the Codex home directory used for auth storage. + fn codex_home(&self) -> PathBuf; + + /// Returns the CLI auth credential storage mode for auth loading. + fn cli_auth_credentials_store_mode(&self) -> AuthCredentialsStoreMode; + + /// Returns the workspace ID that ChatGPT auth should be restricted to, if any. + fn forced_chatgpt_workspace_id(&self) -> Option; +} + impl Debug for AuthManager { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AuthManager") @@ -1404,19 +1421,18 @@ impl AuthManager { )) } - pub fn shared_with_external_auth( - codex_home: PathBuf, + /// Convenience constructor returning an `Arc` wrapper from resolved config. + pub fn shared_from_config( + config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - external_auth: Arc, ) -> Arc { - let manager = Self::shared( - codex_home, + let auth_manager = Self::shared( + config.codex_home(), enable_codex_api_key_env, - auth_credentials_store_mode, + config.cli_auth_credentials_store_mode(), ); - manager.set_external_auth(external_auth); - manager + auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id()); + auth_manager } pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index ad77d8238..247e5a876 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -23,6 +23,7 @@ pub use auth::AuthConfig; pub use auth::AuthCredentialsStoreMode; pub use auth::AuthDotJson; pub use auth::AuthManager; +pub use auth::AuthManagerConfig; pub use auth::CLIENT_ID; pub use auth::CODEX_API_KEY_ENV_VAR; pub use auth::CodexAuth; diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index a0d432117..46470f994 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -56,10 +56,9 @@ impl MessageProcessor { environment_manager: Arc, ) -> Self { let outgoing = Arc::new(outgoing); - let auth_manager = AuthManager::shared( - config.codex_home.clone(), + let auth_manager = AuthManager::shared_from_config( + config.as_ref(), /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, ); let thread_manager = Arc::new(ThreadManager::new( config.as_ref(),