Defer MCP auth discovery until startup failure

This commit is contained in:
Alexey Ldokov
2026-07-03 17:26:13 +01:00
parent da4c8ca57d
commit dbd107a7f3
7 changed files with 162 additions and 72 deletions
@@ -40,6 +40,7 @@ use serde_json::json;
use std::path::Path;
use std::path::PathBuf;
use tempfile::TempDir;
use tokio::net::TcpListener;
use tokio::time::timeout;
use wiremock::Mock;
use wiremock::MockServer;
@@ -1043,6 +1044,45 @@ async fn thread_start_emits_mcp_server_status_updated_notifications() -> Result<
Ok(())
}
#[tokio::test]
async fn thread_start_does_not_wait_for_optional_http_mcp_auth_discovery() -> Result<()> {
let server = create_mock_responses_server_repeating_assistant("Done").await;
let listener = TcpListener::bind("127.0.0.1:0").await?;
let mcp_addr = listener.local_addr()?;
let blackhole_server = tokio::spawn(async move {
let mut connections = Vec::new();
while let Ok((connection, _)) = listener.accept().await {
connections.push(connection);
}
});
let codex_home = TempDir::new()?;
create_config_toml_with_optional_http_mcp(
codex_home.path(),
&server.uri(),
&format!("http://{mcp_addr}/mcp"),
)?;
let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?;
timeout(std::time::Duration::from_secs(30), mcp.initialize()).await??;
let req_id = mcp
.send_thread_start_request_with_auto_env(ThreadStartParams::default())
.await?;
let response: JSONRPCResponse = timeout(
std::time::Duration::from_secs(3),
mcp.read_stream_until_response_message(RequestId::Integer(req_id)),
)
.await
.context("thread/start waited for optional HTTP MCP auth discovery")??;
let response: ThreadStartResponse = to_response(response)?;
assert!(!response.thread.id.is_empty());
blackhole_server.abort();
Ok(())
}
#[tokio::test]
async fn thread_start_surfaces_cloud_config_bundle_load_errors() -> Result<()> {
let server = MockServer::start().await;
@@ -1513,6 +1553,37 @@ stream_max_retries = 0
)
}
fn create_config_toml_with_optional_http_mcp(
codex_home: &Path,
server_uri: &str,
mcp_uri: &str,
) -> std::io::Result<()> {
let config_toml = codex_home.join("config.toml");
std::fs::write(
config_toml,
format!(
r#"
model = "mock-model"
approval_policy = "never"
sandbox_mode = "read-only"
model_provider = "mock_provider"
[model_providers.mock_provider]
name = "Mock provider for test"
base_url = "{server_uri}/v1"
wire_api = "responses"
request_max_retries = 0
stream_max_retries = 0
[mcp_servers.optional_http]
url = "{mcp_uri}"
startup_timeout_sec = 60
"#
),
)
}
#[cfg(target_os = "windows")]
fn broken_mcp_transport_toml() -> &'static str {
r#"command = "cmd"
@@ -22,6 +22,7 @@ use crate::elicitation::ElicitationRequestRouter;
use crate::elicitation::ElicitationReviewerHandle;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::mcp::ToolPluginProvenance;
use crate::mcp::compute_auth_status_entry;
use crate::rmcp_client::AsyncManagedClient;
use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT;
use crate::rmcp_client::MCP_TOOLS_FETCH_UNCACHED_DURATION_METRIC;
@@ -164,6 +165,7 @@ impl McpConnectionManager {
let chatgpt_auth_provider = auth
.filter(|auth| auth.uses_codex_backend())
.map(codex_model_provider::auth_provider_from_auth);
let auth = auth.cloned();
let mcp_servers = mcp_servers.clone();
for (server_name, server) in mcp_servers
.into_iter()
@@ -206,6 +208,7 @@ impl McpConnectionManager {
} else {
chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone())
};
let server_for_auth_status = server.clone();
let async_managed_client = AsyncManagedClient::new(
server_name.clone(),
startup_submit_id.clone(),
@@ -226,11 +229,37 @@ impl McpConnectionManager {
let tx_event = tx_event.clone();
let submit_id = startup_submit_id.clone();
let auth_entry = auth_entries.get(&server_name).cloned();
let auth = auth.clone();
let auth_runtime_context = runtime_context.clone();
join_set.spawn(async move {
let mut outcome = async_managed_client.client().await;
if cancel_token.is_cancelled() {
outcome = Err(StartupOutcomeError::Cancelled);
}
// OAuth discovery may make several network requests. Resolve a missing status only
// after an authentication failure so optional servers cannot delay session startup.
let auth_entry = match (&outcome, auth_entry) {
(Err(error @ StartupOutcomeError::Failed { .. }), None)
if error.is_authentication_required() =>
{
Some(
compute_auth_status_entry(
&server_name,
&server_for_auth_status,
store_mode,
keyring_backend_kind,
auth.as_ref(),
&auth_runtime_context,
)
.await,
)
}
(Err(StartupOutcomeError::Failed { .. }), None) => Some(McpAuthStatusEntry {
config: server_for_auth_status.configured_config().cloned(),
auth_state: McpAuthState::Unsupported,
}),
(_, auth_entry) => auth_entry,
};
let status = match &outcome {
Ok(_) => McpStartupStatus::Ready,
Err(StartupOutcomeError::Cancelled) => McpStartupStatus::Cancelled,
+57 -44
View File
@@ -190,55 +190,68 @@ pub async fn compute_auth_statuses<'a, I>(
where
I: IntoIterator<Item = (&'a String, &'a EffectiveMcpServer)>,
{
let futures = servers.into_iter().map(|(name, server)| {
let name = name.clone();
let config = server.configured_config().cloned();
let runtime_context = runtime_context.clone();
let has_runtime_auth = config
.as_ref()
.is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt))
&& auth.is_some_and(CodexAuth::uses_codex_backend)
&& config.as_ref().is_some_and(|config| {
matches!(
&config.transport,
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var: None,
..
}
)
});
async move {
let auth_state = match config.as_ref() {
Some(config) => {
match compute_auth_status(
&name,
config,
store_mode,
keyring_backend_kind,
has_runtime_auth,
&runtime_context,
)
.await
{
Ok(status) => status,
Err(error) => {
warn!(
"failed to determine auth status for MCP server `{name}`: {error:?}"
);
McpAuthState::Unsupported
}
}
}
None => McpAuthState::Unsupported,
};
let entry = McpAuthStatusEntry { config, auth_state };
(name, entry)
}
let futures = servers.into_iter().map(|(name, server)| async move {
let entry = compute_auth_status_entry(
name,
server,
store_mode,
keyring_backend_kind,
auth,
runtime_context,
)
.await;
(name.clone(), entry)
});
join_all(futures).await.into_iter().collect()
}
pub(crate) async fn compute_auth_status_entry(
name: &str,
server: &EffectiveMcpServer,
store_mode: OAuthCredentialsStoreMode,
keyring_backend_kind: AuthKeyringBackendKind,
auth: Option<&CodexAuth>,
runtime_context: &McpRuntimeContext,
) -> McpAuthStatusEntry {
let config = server.configured_config().cloned();
let has_runtime_auth = config
.as_ref()
.is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt))
&& auth.is_some_and(CodexAuth::uses_codex_backend)
&& config.as_ref().is_some_and(|config| {
matches!(
&config.transport,
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var: None,
..
}
)
});
let auth_state = match config.as_ref() {
Some(config) => {
match compute_auth_status(
name,
config,
store_mode,
keyring_backend_kind,
has_runtime_auth,
runtime_context,
)
.await
{
Ok(status) => status,
Err(error) => {
warn!("failed to determine auth status for MCP server `{name}`: {error:?}");
McpAuthState::Unsupported
}
}
}
None => McpAuthState::Unsupported,
};
McpAuthStatusEntry { config, auth_state }
}
async fn compute_auth_status(
server_name: &str,
config: &McpServerConfig,
+1
View File
@@ -3,6 +3,7 @@ pub use auth::McpOAuthLoginConfig;
pub use auth::McpOAuthLoginSupport;
pub use auth::McpOAuthScopesSource;
pub use auth::ResolvedMcpOAuthScopes;
pub(crate) use auth::compute_auth_status_entry;
pub use auth::compute_auth_statuses;
pub use auth::discover_supported_scopes;
pub use auth::discover_supported_scopes_with_http_client;
+1 -9
View File
@@ -330,14 +330,6 @@ impl Session {
turn_context.cwd.to_path_buf()
});
let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd);
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
auth.as_ref(),
&mcp_runtime_context,
)
.await;
let mcp_startup_cancellation_token = {
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
// The previous runtime owns the old token and may still be serving an in-flight step.
@@ -351,7 +343,7 @@ impl Session {
&mcp_servers,
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
auth_statuses,
HashMap::new(),
&turn_context.approval_policy,
turn_context.sub_id.clone(),
self.get_tx_event(),
-1
View File
@@ -334,7 +334,6 @@ use codex_core_plugins::PluginsManager;
use codex_core_plugins::RecommendedPluginCandidatesInput;
use codex_git_utils::get_git_repo_root;
use codex_mcp::McpConfig;
use codex_mcp::compute_auth_statuses;
use codex_mcp::effective_mcp_servers;
use codex_otel::SessionTelemetry;
use codex_otel::THREAD_STARTED_METRIC;
+3 -18
View File
@@ -676,7 +676,6 @@ impl Session {
.unwrap_or_else(|| session_configuration.cwd().to_path_buf());
let mcp_runtime_context =
McpRuntimeContext::new(Arc::clone(&environment_manager), mcp_runtime_cwd);
let mcp_runtime_context_for_auth = mcp_runtime_context.clone();
let auth_and_mcp_fut = async move {
let auth = auth_manager_clone.auth().await;
let mcp_config = mcp_manager_for_mcp
@@ -689,21 +688,7 @@ impl Session {
.await;
let mcp_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref());
let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config);
let auth_statuses = compute_auth_statuses(
mcp_servers.iter(),
config_for_mcp.mcp_oauth_credentials_store_mode,
config_for_mcp.auth_keyring_backend_kind(),
auth.as_ref(),
&mcp_runtime_context_for_auth,
)
.await;
(
auth,
mcp_config,
mcp_servers,
auth_statuses,
tool_plugin_provenance,
)
(auth, mcp_config, mcp_servers, tool_plugin_provenance)
}
.instrument(info_span!(
"session_init.auth_mcp",
@@ -714,7 +699,7 @@ impl Session {
let (
thread_persistence_result,
state_db_ctx,
(auth, mcp_config, mcp_servers, auth_statuses, tool_plugin_provenance),
(auth, mcp_config, mcp_servers, tool_plugin_provenance),
) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut);
let mut live_thread_init =
@@ -1197,7 +1182,7 @@ impl Session {
&mcp_servers,
config.mcp_oauth_credentials_store_mode,
config.auth_keyring_backend_kind(),
auth_statuses,
HashMap::new(),
&session_configuration.approval_policy,
INITIAL_SUBMIT_ID.to_owned(),
tx_event.clone(),