mirror of
https://github.com/openai/codex.git
synced 2026-10-27 18:48:37 +00:00
Add remote runtime OAuth for MCP
Route Streamable HTTP OAuth discovery, registration, token exchange, and refresh through the selected MCP runtime HTTP client so remote MCP servers do not silently fall back to local network access. Keep browser launch and callback handling local while storing tokens in the existing local MCP OAuth store. Add regression coverage for injected OAuth HTTP requests and stored-token refresh through the exec-server transport. Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -312,6 +312,7 @@ use codex_mcp::McpSnapshotDetail;
|
||||
use codex_mcp::collect_mcp_server_status_snapshot_with_detail;
|
||||
use codex_mcp::discover_supported_scopes;
|
||||
use codex_mcp::effective_mcp_servers;
|
||||
use codex_mcp::http_client_for_server;
|
||||
use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread;
|
||||
use codex_mcp::resolve_oauth_scopes;
|
||||
use codex_model_provider::ProviderAccountError;
|
||||
@@ -355,7 +356,7 @@ use codex_protocol::protocol::USER_MESSAGE_BEGIN;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS;
|
||||
use codex_protocol::user_input::UserInput as CoreInputItem;
|
||||
use codex_rmcp_client::perform_oauth_login_return_url;
|
||||
use codex_rmcp_client::perform_oauth_login_return_url_with_client;
|
||||
use codex_rollout::state_db::StateDbHandle;
|
||||
use codex_rollout::state_db::get_state_db;
|
||||
use codex_rollout::state_db::reconcile_rollout;
|
||||
@@ -5960,15 +5961,36 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
};
|
||||
|
||||
let environment_manager = self.thread_manager.environment_manager();
|
||||
let runtime_environment = match environment_manager.default_environment() {
|
||||
Some(environment) => McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()),
|
||||
None => McpRuntimeEnvironment::new(
|
||||
environment_manager.local_environment(),
|
||||
config.cwd.to_path_buf(),
|
||||
),
|
||||
};
|
||||
let http_client = match http_client_for_server(server, runtime_environment) {
|
||||
Ok(http_client) => http_client,
|
||||
Err(err) => {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: format!("failed to resolve MCP server OAuth environment: {err}"),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let discovered_scopes = if scopes.is_none() && server.scopes.is_none() {
|
||||
discover_supported_scopes(&server.transport).await
|
||||
discover_supported_scopes(&server.transport, http_client.clone()).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let resolved_scopes =
|
||||
resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes);
|
||||
|
||||
match perform_oauth_login_return_url(
|
||||
match perform_oauth_login_return_url_with_client(
|
||||
&name,
|
||||
&url,
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
@@ -5979,6 +6001,7 @@ impl CodexMessageProcessor {
|
||||
timeout_secs,
|
||||
config.mcp_oauth_callback_port,
|
||||
config.mcp_oauth_callback_url.as_deref(),
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -6,10 +6,12 @@ use codex_app_server_protocol::ServerNotification;
|
||||
use codex_config::types::McpServerConfig;
|
||||
use codex_core::config::Config;
|
||||
use codex_mcp::McpOAuthLoginSupport;
|
||||
use codex_mcp::McpRuntimeEnvironment;
|
||||
use codex_mcp::http_client_for_server;
|
||||
use codex_mcp::oauth_login_support;
|
||||
use codex_mcp::resolve_oauth_scopes;
|
||||
use codex_mcp::should_retry_without_scopes;
|
||||
use codex_rmcp_client::perform_oauth_login_silent;
|
||||
use codex_rmcp_client::perform_oauth_login_silent_with_client;
|
||||
use tracing::warn;
|
||||
|
||||
use super::CodexMessageProcessor;
|
||||
@@ -21,7 +23,28 @@ impl CodexMessageProcessor {
|
||||
plugin_mcp_servers: HashMap<String, McpServerConfig>,
|
||||
) {
|
||||
for (name, server) in plugin_mcp_servers {
|
||||
let oauth_config = match oauth_login_support(&server.transport).await {
|
||||
let environment_manager = self.thread_manager.environment_manager();
|
||||
let runtime_environment = match environment_manager.default_environment() {
|
||||
Some(environment) => {
|
||||
McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf())
|
||||
}
|
||||
None => McpRuntimeEnvironment::new(
|
||||
environment_manager.local_environment(),
|
||||
config.cwd.to_path_buf(),
|
||||
),
|
||||
};
|
||||
let http_client = match http_client_for_server(&server, runtime_environment) {
|
||||
Ok(http_client) => http_client,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"failed to resolve MCP OAuth environment for plugin install {name}: {err}"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let oauth_config = match oauth_login_support(&server.transport, http_client.clone())
|
||||
.await
|
||||
{
|
||||
McpOAuthLoginSupport::Supported(config) => config,
|
||||
McpOAuthLoginSupport::Unsupported => continue,
|
||||
McpOAuthLoginSupport::Unknown(err) => {
|
||||
@@ -45,7 +68,7 @@ impl CodexMessageProcessor {
|
||||
let notification_name = name.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let first_attempt = perform_oauth_login_silent(
|
||||
let first_attempt = perform_oauth_login_silent_with_client(
|
||||
&name,
|
||||
&oauth_config.url,
|
||||
store_mode,
|
||||
@@ -55,12 +78,13 @@ impl CodexMessageProcessor {
|
||||
server.oauth_resource.as_deref(),
|
||||
callback_port,
|
||||
callback_url.as_deref(),
|
||||
http_client.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let final_result = match first_attempt {
|
||||
Err(err) if should_retry_without_scopes(&resolved_scopes, &err) => {
|
||||
perform_oauth_login_silent(
|
||||
perform_oauth_login_silent_with_client(
|
||||
&name,
|
||||
&oauth_config.url,
|
||||
store_mode,
|
||||
@@ -70,6 +94,7 @@ impl CodexMessageProcessor {
|
||||
server.oauth_resource.as_deref(),
|
||||
callback_port,
|
||||
callback_url.as_deref(),
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -15,7 +15,10 @@ use codex_core::config::edit::ConfigEditsBuilder;
|
||||
use codex_core::config::find_codex_home;
|
||||
use codex_core::config::load_global_mcp_servers;
|
||||
use codex_core::plugins::PluginsManager;
|
||||
use codex_exec_server::Environment;
|
||||
use codex_exec_server::ReqwestHttpClient;
|
||||
use codex_mcp::McpOAuthLoginSupport;
|
||||
use codex_mcp::McpRuntimeEnvironment;
|
||||
use codex_mcp::ResolvedMcpOAuthScopes;
|
||||
use codex_mcp::compute_auth_statuses;
|
||||
use codex_mcp::discover_supported_scopes;
|
||||
@@ -323,7 +326,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re
|
||||
|
||||
println!("Added global MCP server '{name}'.");
|
||||
|
||||
match oauth_login_support(&transport).await {
|
||||
match oauth_login_support(&transport, Arc::new(ReqwestHttpClient)).await {
|
||||
McpOAuthLoginSupport::Supported(oauth_config) => {
|
||||
println!("Detected OAuth support. Starting OAuth flow…");
|
||||
let resolved_scopes = resolve_oauth_scopes(
|
||||
@@ -417,7 +420,7 @@ async fn run_login(config_overrides: &CliConfigOverrides, login_args: LoginArgs)
|
||||
|
||||
let explicit_scopes = (!scopes.is_empty()).then_some(scopes);
|
||||
let discovered_scopes = if explicit_scopes.is_none() && server.scopes.is_none() {
|
||||
discover_supported_scopes(&server.transport).await
|
||||
discover_supported_scopes(&server.transport, Arc::new(ReqwestHttpClient)).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -490,6 +493,10 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) ->
|
||||
mcp_servers.iter(),
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
/*auth*/ None,
|
||||
McpRuntimeEnvironment::new(
|
||||
Arc::new(Environment::default_for_tests()),
|
||||
config.cwd.to_path_buf(),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
@@ -29,6 +29,7 @@ pub use mcp::McpOAuthScopesSource;
|
||||
pub use mcp::ResolvedMcpOAuthScopes;
|
||||
pub use mcp::compute_auth_statuses;
|
||||
pub use mcp::discover_supported_scopes;
|
||||
pub use mcp::http_client_for_server;
|
||||
pub use mcp::oauth_login_support;
|
||||
pub use mcp::resolve_oauth_scopes;
|
||||
pub use mcp::should_retry_without_scopes;
|
||||
|
||||
@@ -4,14 +4,19 @@ use anyhow::Result;
|
||||
use codex_config::McpServerConfig;
|
||||
use codex_config::McpServerTransportConfig;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::ReqwestHttpClient;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_protocol::protocol::McpAuthStatus;
|
||||
use codex_rmcp_client::OAuthProviderError;
|
||||
use codex_rmcp_client::determine_streamable_http_auth_status;
|
||||
use codex_rmcp_client::discover_streamable_http_oauth;
|
||||
use codex_rmcp_client::determine_streamable_http_auth_status_with_client;
|
||||
use codex_rmcp_client::discover_streamable_http_oauth_with_client;
|
||||
use futures::future::join_all;
|
||||
use std::sync::Arc;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::runtime::McpRuntimeEnvironment;
|
||||
|
||||
use super::CODEX_APPS_MCP_SERVER_NAME;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -49,7 +54,10 @@ pub struct McpAuthStatusEntry {
|
||||
pub auth_status: McpAuthStatus,
|
||||
}
|
||||
|
||||
pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAuthLoginSupport {
|
||||
pub async fn oauth_login_support(
|
||||
transport: &McpServerTransportConfig,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> McpOAuthLoginSupport {
|
||||
let McpServerTransportConfig::StreamableHttp {
|
||||
url,
|
||||
bearer_token_env_var,
|
||||
@@ -64,7 +72,13 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu
|
||||
return McpOAuthLoginSupport::Unsupported;
|
||||
}
|
||||
|
||||
match discover_streamable_http_oauth(url, http_headers.clone(), env_http_headers.clone()).await
|
||||
match discover_streamable_http_oauth_with_client(
|
||||
url,
|
||||
http_headers.clone(),
|
||||
env_http_headers.clone(),
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(discovery)) => McpOAuthLoginSupport::Supported(McpOAuthLoginConfig {
|
||||
url: url.clone(),
|
||||
@@ -79,8 +93,9 @@ pub async fn oauth_login_support(transport: &McpServerTransportConfig) -> McpOAu
|
||||
|
||||
pub async fn discover_supported_scopes(
|
||||
transport: &McpServerTransportConfig,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Option<Vec<String>> {
|
||||
match oauth_login_support(transport).await {
|
||||
match oauth_login_support(transport, http_client).await {
|
||||
McpOAuthLoginSupport::Supported(config) => config.discovered_scopes,
|
||||
McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None,
|
||||
}
|
||||
@@ -129,6 +144,7 @@ pub async fn compute_auth_statuses<'a, I>(
|
||||
servers: I,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
auth: Option<&CodexAuth>,
|
||||
runtime_environment: McpRuntimeEnvironment,
|
||||
) -> HashMap<String, McpAuthStatusEntry>
|
||||
where
|
||||
I: IntoIterator<Item = (&'a String, &'a McpServerConfig)>,
|
||||
@@ -136,6 +152,7 @@ where
|
||||
let futures = servers.into_iter().map(|(name, config)| {
|
||||
let name = name.clone();
|
||||
let config = config.clone();
|
||||
let runtime_environment = runtime_environment.clone();
|
||||
let has_runtime_auth = name == CODEX_APPS_MCP_SERVER_NAME
|
||||
&& auth.is_some_and(CodexAuth::uses_codex_backend)
|
||||
&& matches!(
|
||||
@@ -146,14 +163,21 @@ where
|
||||
}
|
||||
);
|
||||
async move {
|
||||
let auth_status =
|
||||
match compute_auth_status(&name, &config, store_mode, has_runtime_auth).await {
|
||||
Ok(status) => status,
|
||||
Err(error) => {
|
||||
warn!("failed to determine auth status for MCP server `{name}`: {error:?}");
|
||||
McpAuthStatus::Unsupported
|
||||
}
|
||||
};
|
||||
let auth_status = match compute_auth_status(
|
||||
&name,
|
||||
&config,
|
||||
store_mode,
|
||||
has_runtime_auth,
|
||||
runtime_environment,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(status) => status,
|
||||
Err(error) => {
|
||||
warn!("failed to determine auth status for MCP server `{name}`: {error:?}");
|
||||
McpAuthStatus::Unsupported
|
||||
}
|
||||
};
|
||||
let entry = McpAuthStatusEntry {
|
||||
config,
|
||||
auth_status,
|
||||
@@ -170,6 +194,7 @@ async fn compute_auth_status(
|
||||
config: &McpServerConfig,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
has_runtime_auth: bool,
|
||||
runtime_environment: McpRuntimeEnvironment,
|
||||
) -> Result<McpAuthStatus> {
|
||||
if !config.enabled {
|
||||
return Ok(McpAuthStatus::Unsupported);
|
||||
@@ -187,27 +212,57 @@ async fn compute_auth_status(
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
} => {
|
||||
determine_streamable_http_auth_status(
|
||||
let http_client = http_client_for_server(config, runtime_environment)?;
|
||||
determine_streamable_http_auth_status_with_client(
|
||||
server_name,
|
||||
url,
|
||||
bearer_token_env_var.as_deref(),
|
||||
http_headers.clone(),
|
||||
env_http_headers.clone(),
|
||||
store_mode,
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn http_client_for_server(
|
||||
config: &McpServerConfig,
|
||||
runtime_environment: McpRuntimeEnvironment,
|
||||
) -> Result<Arc<dyn HttpClient>> {
|
||||
match config.experimental_environment.as_deref() {
|
||||
None | Some("local") => Ok(Arc::new(ReqwestHttpClient)),
|
||||
Some("remote") => {
|
||||
let environment = runtime_environment.environment();
|
||||
if !environment.is_remote() {
|
||||
anyhow::bail!("remote MCP server requires a remote environment");
|
||||
}
|
||||
Ok(environment.get_http_client())
|
||||
}
|
||||
Some(environment) => anyhow::bail!("unsupported experimental_environment `{environment}`"),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::anyhow;
|
||||
use codex_config::McpServerConfig;
|
||||
use codex_config::McpServerTransportConfig;
|
||||
use codex_exec_server::Environment;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use crate::runtime::McpRuntimeEnvironment;
|
||||
|
||||
use super::McpOAuthScopesSource;
|
||||
use super::OAuthProviderError;
|
||||
use super::ResolvedMcpOAuthScopes;
|
||||
use super::http_client_for_server;
|
||||
use super::resolve_oauth_scopes;
|
||||
use super::should_retry_without_scopes;
|
||||
|
||||
@@ -318,4 +373,71 @@ mod tests {
|
||||
&anyhow!("timed out waiting for OAuth callback"),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_server_uses_local_http_client_even_with_remote_runtime() {
|
||||
let config = streamable_config(/*experimental_environment*/ None);
|
||||
let runtime_environment = remote_runtime_environment();
|
||||
|
||||
assert!(http_client_for_server(&config, runtime_environment).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remote_server_uses_runtime_http_client_when_runtime_is_remote() {
|
||||
let config = streamable_config(Some("remote"));
|
||||
let runtime_environment = remote_runtime_environment();
|
||||
|
||||
assert!(http_client_for_server(&config, runtime_environment).is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_server_without_remote_runtime_returns_clear_error() {
|
||||
let config = streamable_config(Some("remote"));
|
||||
let runtime_environment = McpRuntimeEnvironment::new(
|
||||
Arc::new(Environment::default_for_tests()),
|
||||
PathBuf::from("/tmp"),
|
||||
);
|
||||
|
||||
let Err(error) = http_client_for_server(&config, runtime_environment) else {
|
||||
panic!("remote server should require remote runtime");
|
||||
};
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"remote MCP server requires a remote environment"
|
||||
);
|
||||
}
|
||||
|
||||
fn streamable_config(experimental_environment: Option<&str>) -> McpServerConfig {
|
||||
McpServerConfig {
|
||||
transport: McpServerTransportConfig::StreamableHttp {
|
||||
url: "http://mcp.example.test/mcp".to_string(),
|
||||
bearer_token_env_var: None,
|
||||
http_headers: None,
|
||||
env_http_headers: None,
|
||||
},
|
||||
experimental_environment: experimental_environment.map(str::to_string),
|
||||
enabled: true,
|
||||
required: false,
|
||||
supports_parallel_tool_calls: false,
|
||||
disabled_reason: None,
|
||||
startup_timeout_sec: Some(Duration::from_secs(30)),
|
||||
tool_timeout_sec: None,
|
||||
default_tools_approval_mode: None,
|
||||
enabled_tools: None,
|
||||
disabled_tools: None,
|
||||
scopes: None,
|
||||
oauth_resource: None,
|
||||
tools: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn remote_runtime_environment() -> McpRuntimeEnvironment {
|
||||
McpRuntimeEnvironment::new(
|
||||
Arc::new(
|
||||
Environment::create_for_tests(Some("ws://127.0.0.1:65535".to_string()))
|
||||
.expect("create remote environment"),
|
||||
),
|
||||
PathBuf::from("/tmp"),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ pub use auth::McpOAuthScopesSource;
|
||||
pub use auth::ResolvedMcpOAuthScopes;
|
||||
pub use auth::compute_auth_statuses;
|
||||
pub use auth::discover_supported_scopes;
|
||||
pub use auth::http_client_for_server;
|
||||
pub use auth::oauth_login_support;
|
||||
pub use auth::resolve_oauth_scopes;
|
||||
pub use auth::should_retry_without_scopes;
|
||||
@@ -223,6 +224,7 @@ pub async fn read_mcp_resource(
|
||||
mcp_servers.iter(),
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
auth,
|
||||
runtime_environment.clone(),
|
||||
)
|
||||
.await;
|
||||
let (tx_event, rx_event) = unbounded();
|
||||
@@ -286,6 +288,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
|
||||
mcp_servers.iter(),
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
auth,
|
||||
runtime_environment.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
@@ -253,20 +253,23 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
|
||||
});
|
||||
}
|
||||
|
||||
let environment = environment_manager
|
||||
.default_environment()
|
||||
.unwrap_or_else(|| environment_manager.local_environment());
|
||||
let runtime_environment =
|
||||
McpRuntimeEnvironment::new(environment.clone(), config.cwd.to_path_buf());
|
||||
|
||||
let auth_status_entries = compute_auth_statuses(
|
||||
mcp_servers.iter(),
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
auth.as_ref(),
|
||||
runtime_environment.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let (tx_event, rx_event) = unbounded();
|
||||
drop(rx_event);
|
||||
|
||||
let environment = environment_manager
|
||||
.default_environment()
|
||||
.unwrap_or_else(|| environment_manager.local_environment());
|
||||
|
||||
let (mcp_connection_manager, cancel_token) = McpConnectionManager::new(
|
||||
&mcp_servers,
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
@@ -275,7 +278,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
|
||||
INITIAL_SUBMIT_ID.to_owned(),
|
||||
tx_event,
|
||||
PermissionProfile::default(),
|
||||
McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()),
|
||||
runtime_environment,
|
||||
config.codex_home.to_path_buf(),
|
||||
codex_apps_tools_cache_key(auth.as_ref()),
|
||||
ToolPluginProvenance::default(),
|
||||
|
||||
@@ -11,7 +11,7 @@ use codex_protocol::request_user_input::RequestUserInputArgs;
|
||||
use codex_protocol::request_user_input::RequestUserInputQuestion;
|
||||
use codex_protocol::request_user_input::RequestUserInputQuestionOption;
|
||||
use codex_protocol::request_user_input::RequestUserInputResponse;
|
||||
use codex_rmcp_client::perform_oauth_login;
|
||||
use codex_rmcp_client::perform_oauth_login_with_client;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -20,6 +20,8 @@ use crate::session::session::Session;
|
||||
use crate::session::turn_context::TurnContext;
|
||||
use crate::skills::model::SkillToolDependency;
|
||||
use codex_mcp::McpOAuthLoginSupport;
|
||||
use codex_mcp::McpRuntimeEnvironment;
|
||||
use codex_mcp::http_client_for_server;
|
||||
use codex_mcp::mcp_permission_prompt_is_auto_approved;
|
||||
use codex_mcp::oauth_login_support;
|
||||
use codex_mcp::resolve_oauth_scopes;
|
||||
@@ -126,14 +128,29 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
}
|
||||
|
||||
for (name, server_config) in added {
|
||||
let oauth_config = match oauth_login_support(&server_config.transport).await {
|
||||
McpOAuthLoginSupport::Supported(config) => config,
|
||||
McpOAuthLoginSupport::Unsupported => continue,
|
||||
McpOAuthLoginSupport::Unknown(err) => {
|
||||
warn!("MCP server may or may not require login for dependency {name}: {err}");
|
||||
let runtime_environment = McpRuntimeEnvironment::new(
|
||||
turn_context
|
||||
.environment
|
||||
.clone()
|
||||
.unwrap_or_else(|| sess.services.environment_manager.local_environment()),
|
||||
turn_context.cwd.to_path_buf(),
|
||||
);
|
||||
let http_client = match http_client_for_server(&server_config, runtime_environment) {
|
||||
Ok(http_client) => http_client,
|
||||
Err(err) => {
|
||||
warn!("failed to resolve MCP OAuth environment for dependency {name}: {err}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let oauth_config =
|
||||
match oauth_login_support(&server_config.transport, http_client.clone()).await {
|
||||
McpOAuthLoginSupport::Supported(config) => config,
|
||||
McpOAuthLoginSupport::Unsupported => continue,
|
||||
McpOAuthLoginSupport::Unknown(err) => {
|
||||
warn!("MCP server may or may not require login for dependency {name}: {err}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
sess.notify_background_event(
|
||||
turn_context,
|
||||
@@ -148,7 +165,7 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
server_config.scopes.clone(),
|
||||
oauth_config.discovered_scopes.clone(),
|
||||
);
|
||||
let first_attempt = perform_oauth_login(
|
||||
let first_attempt = perform_oauth_login_with_client(
|
||||
&name,
|
||||
&oauth_config.url,
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
@@ -158,6 +175,7 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
server_config.oauth_resource.as_deref(),
|
||||
config.mcp_oauth_callback_port,
|
||||
config.mcp_oauth_callback_url.as_deref(),
|
||||
http_client.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -171,7 +189,7 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Err(err) = perform_oauth_login(
|
||||
if let Err(err) = perform_oauth_login_with_client(
|
||||
&name,
|
||||
&oauth_config.url,
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
@@ -181,6 +199,7 @@ pub(crate) async fn maybe_install_mcp_dependencies(
|
||||
server_config.oauth_resource.as_deref(),
|
||||
config.mcp_oauth_callback_port,
|
||||
config.mcp_oauth_callback_url.as_deref(),
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -32,6 +32,7 @@ use crate::tasks::UndoTask;
|
||||
use crate::tasks::UserShellCommandMode;
|
||||
use crate::tasks::UserShellCommandTask;
|
||||
use crate::tasks::execute_user_shell_command;
|
||||
use codex_mcp::McpRuntimeEnvironment;
|
||||
use codex_mcp::collect_mcp_snapshot_from_manager;
|
||||
use codex_mcp::compute_auth_statuses;
|
||||
use codex_protocol::models::ContentItem;
|
||||
@@ -538,12 +539,18 @@ pub async fn list_mcp_tools(sess: &Session, config: &Arc<Config>, sub_id: String
|
||||
.mcp_manager
|
||||
.effective_servers(config, auth.as_ref())
|
||||
.await;
|
||||
let environment = sess
|
||||
.services
|
||||
.environment_manager
|
||||
.default_environment()
|
||||
.unwrap_or_else(|| sess.services.environment_manager.local_environment());
|
||||
let snapshot = collect_mcp_snapshot_from_manager(
|
||||
&mcp_connection_manager,
|
||||
compute_auth_statuses(
|
||||
mcp_servers.iter(),
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
auth.as_ref(),
|
||||
McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()),
|
||||
)
|
||||
.await,
|
||||
)
|
||||
|
||||
@@ -219,8 +219,20 @@ impl Session {
|
||||
.tool_plugin_provenance(config.as_ref())
|
||||
.await;
|
||||
let mcp_servers = with_codex_apps_mcp(mcp_servers, auth.as_ref(), &mcp_config);
|
||||
let auth_statuses =
|
||||
compute_auth_statuses(mcp_servers.iter(), store_mode, auth.as_ref()).await;
|
||||
let runtime_environment = McpRuntimeEnvironment::new(
|
||||
turn_context
|
||||
.environment
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.services.environment_manager.local_environment()),
|
||||
turn_context.cwd.to_path_buf(),
|
||||
);
|
||||
let auth_statuses = compute_auth_statuses(
|
||||
mcp_servers.iter(),
|
||||
store_mode,
|
||||
auth.as_ref(),
|
||||
runtime_environment.clone(),
|
||||
)
|
||||
.await;
|
||||
{
|
||||
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
|
||||
guard.cancel();
|
||||
@@ -234,13 +246,7 @@ impl Session {
|
||||
turn_context.sub_id.clone(),
|
||||
self.get_tx_event(),
|
||||
turn_context.permission_profile(),
|
||||
McpRuntimeEnvironment::new(
|
||||
turn_context
|
||||
.environment
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.services.environment_manager.local_environment()),
|
||||
turn_context.cwd.to_path_buf(),
|
||||
),
|
||||
runtime_environment,
|
||||
config.codex_home.to_path_buf(),
|
||||
codex_apps_tools_cache_key(auth.as_ref()),
|
||||
tool_plugin_provenance,
|
||||
|
||||
@@ -445,15 +445,20 @@ impl Session {
|
||||
let auth_manager_clone = Arc::clone(&auth_manager);
|
||||
let config_for_mcp = Arc::clone(&config);
|
||||
let mcp_manager_for_mcp = Arc::clone(&mcp_manager);
|
||||
let environment_manager_for_mcp = Arc::clone(&environment_manager);
|
||||
let auth_and_mcp_fut = async move {
|
||||
let auth = auth_manager_clone.auth().await;
|
||||
let mcp_servers = mcp_manager_for_mcp
|
||||
.effective_servers(&config_for_mcp, auth.as_ref())
|
||||
.await;
|
||||
let environment = environment_manager_for_mcp
|
||||
.default_environment()
|
||||
.unwrap_or_else(|| environment_manager_for_mcp.local_environment());
|
||||
let auth_statuses = compute_auth_statuses(
|
||||
mcp_servers.iter(),
|
||||
config_for_mcp.mcp_oauth_credentials_store_mode,
|
||||
auth.as_ref(),
|
||||
McpRuntimeEnvironment::new(environment, config_for_mcp.cwd.to_path_buf()),
|
||||
)
|
||||
.await;
|
||||
(auth, mcp_servers, auth_statuses)
|
||||
|
||||
@@ -1,31 +1,19 @@
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Error;
|
||||
use anyhow::Result;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::ReqwestHttpClient;
|
||||
use codex_protocol::protocol::McpAuthStatus;
|
||||
use reqwest::Client;
|
||||
use reqwest::StatusCode;
|
||||
use reqwest::Url;
|
||||
use reqwest::header::AUTHORIZATION;
|
||||
use reqwest::header::HeaderMap;
|
||||
use serde::Deserialize;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::mcp_oauth_http::OAuthHttpClient;
|
||||
use crate::mcp_oauth_http::StreamableHttpOAuthDiscovery;
|
||||
use crate::oauth::has_oauth_tokens;
|
||||
use crate::utils::apply_default_headers;
|
||||
use crate::utils::build_default_headers;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
|
||||
const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const OAUTH_DISCOVERY_HEADER: &str = "MCP-Protocol-Version";
|
||||
const OAUTH_DISCOVERY_VERSION: &str = "2024-11-05";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StreamableHttpOAuthDiscovery {
|
||||
pub scopes_supported: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// Determine the authentication status for a streamable HTTP MCP server.
|
||||
pub async fn determine_streamable_http_auth_status(
|
||||
server_name: &str,
|
||||
@@ -34,12 +22,36 @@ pub async fn determine_streamable_http_auth_status(
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
) -> Result<McpAuthStatus> {
|
||||
determine_streamable_http_auth_status_with_client(
|
||||
server_name,
|
||||
url,
|
||||
bearer_token_env_var,
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
store_mode,
|
||||
Arc::new(ReqwestHttpClient),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Determine the authentication status for a streamable HTTP MCP server using
|
||||
/// the caller-selected runtime HTTP client.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn determine_streamable_http_auth_status_with_client(
|
||||
server_name: &str,
|
||||
url: &str,
|
||||
bearer_token_env_var: Option<&str>,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<McpAuthStatus> {
|
||||
if bearer_token_env_var.is_some() {
|
||||
return Ok(McpAuthStatus::BearerToken);
|
||||
}
|
||||
|
||||
let default_headers = build_default_headers(http_headers, env_http_headers)?;
|
||||
let default_headers = build_default_headers(http_headers.clone(), env_http_headers.clone())?;
|
||||
if default_headers.contains_key(AUTHORIZATION) {
|
||||
return Ok(McpAuthStatus::BearerToken);
|
||||
}
|
||||
@@ -48,7 +60,8 @@ pub async fn determine_streamable_http_auth_status(
|
||||
return Ok(McpAuthStatus::OAuth);
|
||||
}
|
||||
|
||||
match discover_streamable_http_oauth_with_headers(url, &default_headers).await {
|
||||
let oauth_http = OAuthHttpClient::from_default_headers(http_client, default_headers);
|
||||
match oauth_http.discover(url).await {
|
||||
Ok(Some(_)) => Ok(McpAuthStatus::NotLoggedIn),
|
||||
Ok(None) => Ok(McpAuthStatus::Unsupported),
|
||||
Err(error) => {
|
||||
@@ -74,121 +87,23 @@ pub async fn discover_streamable_http_oauth(
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
) -> Result<Option<StreamableHttpOAuthDiscovery>> {
|
||||
let default_headers = build_default_headers(http_headers, env_http_headers)?;
|
||||
discover_streamable_http_oauth_with_headers(url, &default_headers).await
|
||||
discover_streamable_http_oauth_with_client(
|
||||
url,
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
Arc::new(ReqwestHttpClient),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn discover_streamable_http_oauth_with_headers(
|
||||
pub async fn discover_streamable_http_oauth_with_client(
|
||||
url: &str,
|
||||
default_headers: &HeaderMap,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<Option<StreamableHttpOAuthDiscovery>> {
|
||||
let base_url = Url::parse(url)?;
|
||||
|
||||
// Use no_proxy to avoid a bug in the system-configuration crate that
|
||||
// can result in a panic. See #8912.
|
||||
let builder = Client::builder().timeout(DISCOVERY_TIMEOUT).no_proxy();
|
||||
let client = apply_default_headers(builder, default_headers).build()?;
|
||||
|
||||
let mut last_error: Option<Error> = None;
|
||||
for candidate_path in discovery_paths(base_url.path()) {
|
||||
let mut discovery_url = base_url.clone();
|
||||
discovery_url.set_path(&candidate_path);
|
||||
|
||||
let response = match client
|
||||
.get(discovery_url.clone())
|
||||
.header(OAUTH_DISCOVERY_HEADER, OAUTH_DISCOVERY_VERSION)
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(err) => {
|
||||
last_error = Some(err.into());
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if response.status() != StatusCode::OK {
|
||||
continue;
|
||||
}
|
||||
|
||||
let metadata = match response.json::<OAuthDiscoveryMetadata>().await {
|
||||
Ok(metadata) => metadata,
|
||||
Err(err) => {
|
||||
last_error = Some(err.into());
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if metadata.authorization_endpoint.is_some() && metadata.token_endpoint.is_some() {
|
||||
return Ok(Some(StreamableHttpOAuthDiscovery {
|
||||
scopes_supported: normalize_scopes(metadata.scopes_supported),
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(err) = last_error {
|
||||
debug!("OAuth discovery requests failed for {url}: {err:?}");
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OAuthDiscoveryMetadata {
|
||||
#[serde(default)]
|
||||
authorization_endpoint: Option<String>,
|
||||
#[serde(default)]
|
||||
token_endpoint: Option<String>,
|
||||
#[serde(default)]
|
||||
scopes_supported: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
fn normalize_scopes(scopes_supported: Option<Vec<String>>) -> Option<Vec<String>> {
|
||||
let scopes_supported = scopes_supported?;
|
||||
|
||||
let mut normalized = Vec::new();
|
||||
for scope in scopes_supported {
|
||||
let scope = scope.trim();
|
||||
if scope.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let scope = scope.to_string();
|
||||
if !normalized.contains(&scope) {
|
||||
normalized.push(scope);
|
||||
}
|
||||
}
|
||||
|
||||
if normalized.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(normalized)
|
||||
}
|
||||
}
|
||||
|
||||
/// Implements RFC 8414 section 3.1 for discovering well-known oauth endpoints.
|
||||
/// This is a requirement for MCP servers to support OAuth.
|
||||
/// https://datatracker.ietf.org/doc/html/rfc8414#section-3.1
|
||||
/// https://github.com/modelcontextprotocol/rust-sdk/blob/main/crates/rmcp/src/transport/auth.rs#L182
|
||||
fn discovery_paths(base_path: &str) -> Vec<String> {
|
||||
let trimmed = base_path.trim_start_matches('/').trim_end_matches('/');
|
||||
let canonical = "/.well-known/oauth-authorization-server".to_string();
|
||||
|
||||
if trimmed.is_empty() {
|
||||
return vec![canonical];
|
||||
}
|
||||
|
||||
let mut candidates = Vec::new();
|
||||
let mut push_unique = |candidate: String| {
|
||||
if !candidates.contains(&candidate) {
|
||||
candidates.push(candidate);
|
||||
}
|
||||
};
|
||||
|
||||
push_unique(format!("{canonical}/{trimmed}"));
|
||||
push_unique(format!("/{trimmed}/.well-known/oauth-authorization-server"));
|
||||
push_unique(canonical);
|
||||
|
||||
candidates
|
||||
let oauth_http = OAuthHttpClient::new(http_client, http_headers, env_http_headers)?;
|
||||
oauth_http.discover(url).await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -8,6 +8,7 @@ use std::time::Duration;
|
||||
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::Json;
|
||||
use axum::extract::State;
|
||||
use axum::http::HeaderMap;
|
||||
@@ -122,6 +123,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
SESSION_POST_FAILURE_CONTROL_PATH,
|
||||
post(arm_session_post_failure),
|
||||
)
|
||||
.route("/oauth/token", post(oauth_token))
|
||||
.route(
|
||||
"/.well-known/oauth-authorization-server/mcp",
|
||||
get({
|
||||
@@ -139,6 +141,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
serde_json::to_vec(&json!({
|
||||
"authorization_endpoint": format!("{metadata_base}/oauth/authorize"),
|
||||
"token_endpoint": format!("{metadata_base}/oauth/token"),
|
||||
"registration_endpoint": format!("{metadata_base}/oauth/register"),
|
||||
"scopes_supported": [""],
|
||||
})).expect("failed to serialize metadata"),
|
||||
))
|
||||
@@ -386,7 +389,8 @@ async fn require_bearer(
|
||||
request: Request<Body>,
|
||||
next: Next,
|
||||
) -> Result<Response, StatusCode> {
|
||||
if request.uri().path().contains("/.well-known/") {
|
||||
let path = request.uri().path();
|
||||
if path.contains("/.well-known/") || path.starts_with("/oauth/") {
|
||||
return Ok(next.run(request).await);
|
||||
}
|
||||
if request
|
||||
@@ -400,6 +404,33 @@ async fn require_bearer(
|
||||
}
|
||||
}
|
||||
|
||||
async fn oauth_token(body: Bytes) -> Result<Response, StatusCode> {
|
||||
let form = String::from_utf8(body.to_vec()).map_err(|_| StatusCode::BAD_REQUEST)?;
|
||||
if !form.contains("grant_type=refresh_token") {
|
||||
return Err(StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
let access_token =
|
||||
std::env::var("MCP_OAUTH_ACCESS_TOKEN").unwrap_or_else(|_| "refreshed-oauth-token".into());
|
||||
let refresh_token =
|
||||
std::env::var("MCP_OAUTH_REFRESH_TOKEN").unwrap_or_else(|_| "refresh-token".into());
|
||||
|
||||
#[expect(clippy::expect_used)]
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.body(Body::from(
|
||||
serde_json::to_vec(&json!({
|
||||
"access_token": access_token,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"refresh_token": refresh_token,
|
||||
}))
|
||||
.expect("failed to serialize token response"),
|
||||
))
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
|
||||
}
|
||||
|
||||
async fn arm_session_post_failure(
|
||||
State(state): State<SessionFailureState>,
|
||||
Json(request): Json<ArmSessionPostFailureRequest>,
|
||||
|
||||
@@ -35,6 +35,8 @@ use rmcp::transport::streamable_http_client::StreamableHttpPostResponse;
|
||||
use sse_stream::Sse;
|
||||
use sse_stream::SseStream;
|
||||
|
||||
use crate::oauth::OAuthPersistor;
|
||||
|
||||
const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
|
||||
const JSON_MIME_TYPE: &str = "application/json";
|
||||
const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
|
||||
@@ -45,6 +47,7 @@ pub(crate) struct StreamableHttpClientAdapter {
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
default_headers: HeaderMap,
|
||||
auth_provider: Option<SharedAuthProvider>,
|
||||
oauth_persistor: Option<OAuthPersistor>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
@@ -55,6 +58,8 @@ pub(crate) enum StreamableHttpClientAdapterError {
|
||||
HttpRequest(#[from] ExecServerError),
|
||||
#[error("invalid HTTP header: {0}")]
|
||||
Header(String),
|
||||
#[error("OAuth token error: {0}")]
|
||||
Auth(String),
|
||||
}
|
||||
|
||||
impl StreamableHttpClientAdapter {
|
||||
@@ -62,11 +67,13 @@ impl StreamableHttpClientAdapter {
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
default_headers: HeaderMap,
|
||||
auth_provider: Option<SharedAuthProvider>,
|
||||
oauth_persistor: Option<OAuthPersistor>,
|
||||
) -> Self {
|
||||
Self {
|
||||
http_client,
|
||||
default_headers,
|
||||
auth_provider,
|
||||
oauth_persistor,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -82,7 +89,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
let mut headers = self.default_headers.clone();
|
||||
self.add_auth_headers(&mut headers);
|
||||
self.add_auth_headers(&mut headers).await?;
|
||||
insert_header(
|
||||
&mut headers,
|
||||
ACCEPT,
|
||||
@@ -179,7 +186,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
auth_token: Option<String>,
|
||||
) -> std::result::Result<(), StreamableHttpError<Self::Error>> {
|
||||
let mut headers = self.default_headers.clone();
|
||||
self.add_auth_headers(&mut headers);
|
||||
self.add_auth_headers(&mut headers).await?;
|
||||
if let Some(auth_token) = auth_token {
|
||||
insert_header(
|
||||
&mut headers,
|
||||
@@ -232,7 +239,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
let mut headers = self.default_headers.clone();
|
||||
self.add_auth_headers(&mut headers);
|
||||
self.add_auth_headers(&mut headers).await?;
|
||||
insert_header(
|
||||
&mut headers,
|
||||
ACCEPT,
|
||||
@@ -308,10 +315,27 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
}
|
||||
|
||||
impl StreamableHttpClientAdapter {
|
||||
fn add_auth_headers(&self, headers: &mut HeaderMap) {
|
||||
async fn add_auth_headers(
|
||||
&self,
|
||||
headers: &mut HeaderMap,
|
||||
) -> std::result::Result<(), StreamableHttpError<StreamableHttpClientAdapterError>> {
|
||||
if let Some(auth_provider) = &self.auth_provider {
|
||||
headers.extend(auth_provider.to_auth_headers());
|
||||
}
|
||||
if !headers.contains_key(AUTHORIZATION)
|
||||
&& let Some(oauth_persistor) = &self.oauth_persistor
|
||||
&& let Some(access_token) = oauth_persistor.access_token().await.map_err(|err| {
|
||||
StreamableHttpError::Client(StreamableHttpClientAdapterError::Auth(err.to_string()))
|
||||
})?
|
||||
{
|
||||
insert_header(
|
||||
headers,
|
||||
AUTHORIZATION,
|
||||
format!("Bearer {access_token}"),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ mod elicitation_client_service;
|
||||
mod executor_process_transport;
|
||||
mod http_client_adapter;
|
||||
mod logging_client_handler;
|
||||
mod mcp_oauth_http;
|
||||
mod oauth;
|
||||
mod perform_oauth_login;
|
||||
mod program_resolver;
|
||||
@@ -10,11 +11,13 @@ mod rmcp_client;
|
||||
mod stdio_server_launcher;
|
||||
mod utils;
|
||||
|
||||
pub use auth_status::StreamableHttpOAuthDiscovery;
|
||||
pub use auth_status::determine_streamable_http_auth_status;
|
||||
pub use auth_status::determine_streamable_http_auth_status_with_client;
|
||||
pub use auth_status::discover_streamable_http_oauth;
|
||||
pub use auth_status::discover_streamable_http_oauth_with_client;
|
||||
pub use auth_status::supports_oauth_login;
|
||||
pub use codex_protocol::protocol::McpAuthStatus;
|
||||
pub use mcp_oauth_http::StreamableHttpOAuthDiscovery;
|
||||
pub use oauth::StoredOAuthTokens;
|
||||
pub use oauth::WrappedOAuthTokenResponse;
|
||||
pub use oauth::delete_oauth_tokens;
|
||||
@@ -24,7 +27,10 @@ pub use perform_oauth_login::OAuthProviderError;
|
||||
pub use perform_oauth_login::OauthLoginHandle;
|
||||
pub use perform_oauth_login::perform_oauth_login;
|
||||
pub use perform_oauth_login::perform_oauth_login_return_url;
|
||||
pub use perform_oauth_login::perform_oauth_login_return_url_with_client;
|
||||
pub use perform_oauth_login::perform_oauth_login_silent;
|
||||
pub use perform_oauth_login::perform_oauth_login_silent_with_client;
|
||||
pub use perform_oauth_login::perform_oauth_login_with_client;
|
||||
pub use rmcp::model::ElicitationAction;
|
||||
pub use rmcp_client::Elicitation;
|
||||
pub use rmcp_client::ElicitationResponse;
|
||||
|
||||
@@ -0,0 +1,577 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::HttpHeader;
|
||||
use codex_exec_server::HttpRequestParams;
|
||||
use oauth2::AuthUrl;
|
||||
use oauth2::AuthorizationCode;
|
||||
use oauth2::ClientId;
|
||||
use oauth2::ClientSecret;
|
||||
use oauth2::CsrfToken;
|
||||
use oauth2::EmptyExtraTokenFields;
|
||||
use oauth2::PkceCodeChallenge;
|
||||
use oauth2::PkceCodeVerifier;
|
||||
use oauth2::RedirectUrl;
|
||||
use oauth2::RefreshToken;
|
||||
use oauth2::RequestTokenError;
|
||||
use oauth2::Scope;
|
||||
use oauth2::StandardErrorResponse;
|
||||
use oauth2::StandardTokenResponse;
|
||||
use oauth2::TokenResponse;
|
||||
use oauth2::TokenUrl;
|
||||
use oauth2::basic::BasicClient;
|
||||
use oauth2::basic::BasicErrorResponseType;
|
||||
use oauth2::basic::BasicTokenType;
|
||||
use reqwest::StatusCode;
|
||||
use reqwest::Url;
|
||||
use reqwest::header::HeaderMap;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::WrappedOAuthTokenResponse;
|
||||
use crate::oauth::StoredOAuthTokens;
|
||||
use crate::oauth::compute_expires_at_millis;
|
||||
use crate::utils::build_default_headers;
|
||||
|
||||
const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
const OAUTH_HTTP_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const OAUTH_DISCOVERY_HEADER: &str = "MCP-Protocol-Version";
|
||||
const OAUTH_DISCOVERY_VERSION: &str = "2024-11-05";
|
||||
|
||||
type OAuthTokenResponse = StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>;
|
||||
type OAuthErrorResponse = StandardErrorResponse<BasicErrorResponseType>;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StreamableHttpOAuthDiscovery {
|
||||
pub scopes_supported: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct OAuthHttpClient {
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
default_headers: HeaderMap,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("{0}")]
|
||||
pub(crate) struct OAuthHttpError(String);
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct OAuthAuthorizationSession {
|
||||
pub authorization_url: String,
|
||||
pub csrf_state: CsrfToken,
|
||||
pkce_verifier: PkceCodeVerifier,
|
||||
client: BasicClient<
|
||||
oauth2::EndpointSet,
|
||||
oauth2::EndpointNotSet,
|
||||
oauth2::EndpointNotSet,
|
||||
oauth2::EndpointNotSet,
|
||||
oauth2::EndpointSet,
|
||||
>,
|
||||
client_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct OAuthClientConfig {
|
||||
client_id: String,
|
||||
client_secret: Option<String>,
|
||||
}
|
||||
|
||||
impl OAuthHttpClient {
|
||||
pub(crate) fn new(
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
) -> Result<Self> {
|
||||
let default_headers = build_default_headers(http_headers, env_http_headers)?;
|
||||
Ok(Self {
|
||||
http_client,
|
||||
default_headers,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn from_default_headers(
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
default_headers: HeaderMap,
|
||||
) -> Self {
|
||||
Self {
|
||||
http_client,
|
||||
default_headers,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn discover(&self, url: &str) -> Result<Option<StreamableHttpOAuthDiscovery>> {
|
||||
let metadata = match self.discover_metadata(url).await? {
|
||||
Some(metadata) => metadata,
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
Ok(Some(StreamableHttpOAuthDiscovery {
|
||||
scopes_supported: normalize_scopes(metadata.scopes_supported),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn discover_metadata(&self, url: &str) -> Result<Option<OAuthDiscoveryMetadata>> {
|
||||
let base_url = Url::parse(url)?;
|
||||
let mut last_error: Option<anyhow::Error> = None;
|
||||
|
||||
for candidate_path in discovery_paths(base_url.path()) {
|
||||
let mut discovery_url = base_url.clone();
|
||||
discovery_url.set_path(&candidate_path);
|
||||
|
||||
let response = match self
|
||||
.request(
|
||||
"GET",
|
||||
discovery_url.as_str(),
|
||||
vec![HttpHeader {
|
||||
name: OAUTH_DISCOVERY_HEADER.to_string(),
|
||||
value: OAUTH_DISCOVERY_VERSION.to_string(),
|
||||
}],
|
||||
None,
|
||||
Some(DISCOVERY_TIMEOUT),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(err) => {
|
||||
last_error = Some(err);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if response.status != StatusCode::OK.as_u16() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let metadata = match serde_json::from_slice::<OAuthDiscoveryMetadata>(&response.body.0)
|
||||
{
|
||||
Ok(metadata) => metadata,
|
||||
Err(err) => {
|
||||
last_error = Some(err.into());
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if metadata.authorization_endpoint.is_some() && metadata.token_endpoint.is_some() {
|
||||
return Ok(Some(metadata));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(err) = last_error {
|
||||
debug!("OAuth discovery requests failed for {url}: {err:?}");
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
pub(crate) async fn start_authorization(
|
||||
&self,
|
||||
server_url: &str,
|
||||
scopes: &[String],
|
||||
redirect_uri: &str,
|
||||
client_name: &str,
|
||||
) -> Result<OAuthAuthorizationSession> {
|
||||
let metadata = self
|
||||
.discover_metadata(server_url)
|
||||
.await?
|
||||
.ok_or_else(|| anyhow!("No authorization support detected"))?;
|
||||
let client_config = self
|
||||
.register_client(&metadata, client_name, redirect_uri)
|
||||
.await?;
|
||||
let client = oauth_client(&metadata, &client_config, redirect_uri)?;
|
||||
let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256();
|
||||
let mut auth_request = client
|
||||
.authorize_url(CsrfToken::new_random)
|
||||
.set_pkce_challenge(pkce_challenge);
|
||||
for scope in scopes {
|
||||
auth_request = auth_request.add_scope(Scope::new(scope.clone()));
|
||||
}
|
||||
let (authorization_url, csrf_state) = auth_request.url();
|
||||
|
||||
Ok(OAuthAuthorizationSession {
|
||||
authorization_url: authorization_url.to_string(),
|
||||
csrf_state,
|
||||
pkce_verifier,
|
||||
client,
|
||||
client_id: client_config.client_id,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn exchange_code(
|
||||
&self,
|
||||
session: OAuthAuthorizationSession,
|
||||
code: &str,
|
||||
csrf_state: &str,
|
||||
) -> Result<(String, OAuthTokenResponse)> {
|
||||
if session.csrf_state.secret() != csrf_state {
|
||||
return Err(anyhow!(
|
||||
"OAuth callback state did not match authorization request"
|
||||
));
|
||||
}
|
||||
|
||||
let http_client = self.clone();
|
||||
let token = session
|
||||
.client
|
||||
.exchange_code(AuthorizationCode::new(code.to_string()))
|
||||
.set_pkce_verifier(session.pkce_verifier)
|
||||
.request_async(&|request| {
|
||||
let http_client = http_client.clone();
|
||||
async move { http_client.oauth_request(request).await }
|
||||
})
|
||||
.await
|
||||
.or_else(parse_token_from_parse_error)?;
|
||||
|
||||
Ok((session.client_id, token))
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_token(
|
||||
&self,
|
||||
tokens: &StoredOAuthTokens,
|
||||
) -> Result<Option<StoredOAuthTokens>> {
|
||||
let refresh_token = match tokens.token_response.0.refresh_token() {
|
||||
Some(refresh_token) => refresh_token.secret().to_string(),
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let metadata = self
|
||||
.discover_metadata(&tokens.url)
|
||||
.await?
|
||||
.ok_or_else(|| anyhow!("No authorization support detected"))?;
|
||||
let client_config = OAuthClientConfig {
|
||||
client_id: tokens.client_id.clone(),
|
||||
client_secret: None,
|
||||
};
|
||||
let client = oauth_client(&metadata, &client_config, &tokens.url)?;
|
||||
let http_client = self.clone();
|
||||
let token = client
|
||||
.exchange_refresh_token(&RefreshToken::new(refresh_token))
|
||||
.request_async(&|request| {
|
||||
let http_client = http_client.clone();
|
||||
async move { http_client.oauth_request(request).await }
|
||||
})
|
||||
.await
|
||||
.or_else(parse_token_from_parse_error)?;
|
||||
|
||||
let expires_at = compute_expires_at_millis(&token);
|
||||
Ok(Some(StoredOAuthTokens {
|
||||
server_name: tokens.server_name.clone(),
|
||||
url: tokens.url.clone(),
|
||||
client_id: tokens.client_id.clone(),
|
||||
token_response: WrappedOAuthTokenResponse(token),
|
||||
expires_at,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn register_client(
|
||||
&self,
|
||||
metadata: &OAuthDiscoveryMetadata,
|
||||
client_name: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Result<OAuthClientConfig> {
|
||||
let registration_url = metadata
|
||||
.registration_endpoint
|
||||
.as_deref()
|
||||
.ok_or_else(|| anyhow!("Dynamic client registration not supported"))?;
|
||||
|
||||
if let Some(response_types_supported) = metadata.response_types_supported.as_ref()
|
||||
&& !response_types_supported.iter().any(|value| value == "code")
|
||||
{
|
||||
return Err(anyhow!(
|
||||
"OAuth server does not support authorization code flow"
|
||||
));
|
||||
}
|
||||
|
||||
let body = serde_json::to_vec(&ClientRegistrationRequest {
|
||||
client_name: client_name.to_string(),
|
||||
redirect_uris: vec![redirect_uri.to_string()],
|
||||
grant_types: vec![
|
||||
"authorization_code".to_string(),
|
||||
"refresh_token".to_string(),
|
||||
],
|
||||
token_endpoint_auth_method: "none".to_string(),
|
||||
response_types: vec!["code".to_string()],
|
||||
})?;
|
||||
|
||||
let response = self
|
||||
.request(
|
||||
"POST",
|
||||
registration_url,
|
||||
vec![HttpHeader {
|
||||
name: reqwest::header::CONTENT_TYPE.to_string(),
|
||||
value: "application/json".to_string(),
|
||||
}],
|
||||
Some(body),
|
||||
Some(OAUTH_HTTP_TIMEOUT),
|
||||
)
|
||||
.await?;
|
||||
|
||||
if !status_is_success(response.status) {
|
||||
return Err(anyhow!(
|
||||
"Dynamic registration failed: HTTP {}: {}",
|
||||
response.status,
|
||||
String::from_utf8_lossy(&response.body.0)
|
||||
));
|
||||
}
|
||||
|
||||
let registration: ClientRegistrationResponse = serde_json::from_slice(&response.body.0)
|
||||
.context("failed to parse registration response")?;
|
||||
Ok(OAuthClientConfig {
|
||||
client_id: registration.client_id,
|
||||
client_secret: registration
|
||||
.client_secret
|
||||
.filter(|secret| !secret.is_empty()),
|
||||
})
|
||||
}
|
||||
|
||||
async fn oauth_request(
|
||||
&self,
|
||||
request: oauth2::HttpRequest,
|
||||
) -> std::result::Result<oauth2::HttpResponse, OAuthHttpError> {
|
||||
let (parts, body) = request.into_parts();
|
||||
let headers = parts
|
||||
.headers
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
let value = value
|
||||
.to_str()
|
||||
.map_err(|err| OAuthHttpError(format!("invalid OAuth header value: {err}")))?;
|
||||
Ok(HttpHeader {
|
||||
name: name.to_string(),
|
||||
value: value.to_string(),
|
||||
})
|
||||
})
|
||||
.collect::<std::result::Result<Vec<_>, OAuthHttpError>>()?;
|
||||
let response = self
|
||||
.request(
|
||||
parts.method.as_str(),
|
||||
parts.uri.to_string().as_str(),
|
||||
headers,
|
||||
Some(body),
|
||||
Some(OAUTH_HTTP_TIMEOUT),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| OAuthHttpError(err.to_string()))?;
|
||||
|
||||
let mut oauth_response = oauth2::HttpResponse::new(response.body.0);
|
||||
*oauth_response.status_mut() = oauth2::http::StatusCode::from_u16(response.status)
|
||||
.map_err(|err| OAuthHttpError(format!("invalid OAuth response status: {err}")))?;
|
||||
for header in response.headers {
|
||||
let name = oauth2::http::HeaderName::from_bytes(header.name.as_bytes())
|
||||
.map_err(|err| OAuthHttpError(format!("invalid OAuth response header: {err}")))?;
|
||||
let value = oauth2::http::HeaderValue::from_str(&header.value).map_err(|err| {
|
||||
OAuthHttpError(format!("invalid OAuth response header value: {err}"))
|
||||
})?;
|
||||
oauth_response.headers_mut().append(name, value);
|
||||
}
|
||||
Ok(oauth_response)
|
||||
}
|
||||
|
||||
async fn request(
|
||||
&self,
|
||||
method: &str,
|
||||
url: &str,
|
||||
extra_headers: Vec<HttpHeader>,
|
||||
body: Option<Vec<u8>>,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<codex_exec_server::HttpRequestResponse> {
|
||||
let mut headers = protocol_headers(&self.default_headers)?;
|
||||
headers.extend(extra_headers);
|
||||
self.http_client
|
||||
.http_request(HttpRequestParams {
|
||||
method: method.to_string(),
|
||||
url: url.to_string(),
|
||||
headers,
|
||||
body: body.map(Into::into),
|
||||
timeout_ms: timeout
|
||||
.map(|timeout| timeout.as_millis().clamp(1, u64::MAX as u128) as u64),
|
||||
request_id: "oauth-request".to_string(),
|
||||
stream_response: false,
|
||||
})
|
||||
.await
|
||||
.map_err(|err| anyhow!(err))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OAuthDiscoveryMetadata {
|
||||
#[serde(default)]
|
||||
authorization_endpoint: Option<String>,
|
||||
#[serde(default)]
|
||||
token_endpoint: Option<String>,
|
||||
#[serde(default)]
|
||||
registration_endpoint: Option<String>,
|
||||
#[serde(default)]
|
||||
scopes_supported: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
response_types_supported: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ClientRegistrationRequest {
|
||||
client_name: String,
|
||||
redirect_uris: Vec<String>,
|
||||
grant_types: Vec<String>,
|
||||
token_endpoint_auth_method: String,
|
||||
response_types: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ClientRegistrationResponse {
|
||||
client_id: String,
|
||||
client_secret: Option<String>,
|
||||
}
|
||||
|
||||
fn oauth_client(
|
||||
metadata: &OAuthDiscoveryMetadata,
|
||||
config: &OAuthClientConfig,
|
||||
redirect_uri: &str,
|
||||
) -> Result<
|
||||
BasicClient<
|
||||
oauth2::EndpointSet,
|
||||
oauth2::EndpointNotSet,
|
||||
oauth2::EndpointNotSet,
|
||||
oauth2::EndpointNotSet,
|
||||
oauth2::EndpointSet,
|
||||
>,
|
||||
> {
|
||||
let authorization_endpoint = metadata
|
||||
.authorization_endpoint
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow!("OAuth metadata did not include authorization endpoint"))?;
|
||||
let token_endpoint = metadata
|
||||
.token_endpoint
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow!("OAuth metadata did not include token endpoint"))?;
|
||||
let mut client = BasicClient::new(ClientId::new(config.client_id.clone()))
|
||||
.set_auth_uri(AuthUrl::new(authorization_endpoint)?)
|
||||
.set_token_uri(TokenUrl::new(token_endpoint)?)
|
||||
.set_redirect_uri(RedirectUrl::new(redirect_uri.to_string())?);
|
||||
if let Some(secret) = config.client_secret.clone() {
|
||||
client = client.set_client_secret(ClientSecret::new(secret));
|
||||
}
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
fn parse_token_from_parse_error(
|
||||
error: RequestTokenError<OAuthHttpError, OAuthErrorResponse>,
|
||||
) -> std::result::Result<OAuthTokenResponse, RequestTokenError<OAuthHttpError, OAuthErrorResponse>>
|
||||
{
|
||||
match error {
|
||||
RequestTokenError::Parse(parse_error, body) => {
|
||||
match serde_json::from_slice::<OAuthTokenResponse>(&body) {
|
||||
Ok(parsed) => Ok(parsed),
|
||||
Err(_) => Err(RequestTokenError::Parse(parse_error, body)),
|
||||
}
|
||||
}
|
||||
error => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn protocol_headers(headers: &HeaderMap) -> Result<Vec<HttpHeader>> {
|
||||
headers
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
let value = value
|
||||
.to_str()
|
||||
.with_context(|| format!("invalid HTTP header value for `{name}`"))?;
|
||||
Ok(HttpHeader {
|
||||
name: name.to_string(),
|
||||
value: value.to_string(),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn status_is_success(status: u16) -> bool {
|
||||
(200..300).contains(&status)
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_scopes(scopes_supported: Option<Vec<String>>) -> Option<Vec<String>> {
|
||||
let scopes_supported = scopes_supported?;
|
||||
|
||||
let mut normalized = Vec::new();
|
||||
for scope in scopes_supported {
|
||||
let scope = scope.trim();
|
||||
if scope.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let scope = scope.to_string();
|
||||
if !normalized.contains(&scope) {
|
||||
normalized.push(scope);
|
||||
}
|
||||
}
|
||||
|
||||
if normalized.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(normalized)
|
||||
}
|
||||
}
|
||||
|
||||
/// Implements RFC 8414 section 3.1 for discovering well-known oauth endpoints.
|
||||
pub(crate) fn discovery_paths(base_path: &str) -> Vec<String> {
|
||||
let trimmed = base_path.trim_start_matches('/').trim_end_matches('/');
|
||||
let canonical = "/.well-known/oauth-authorization-server".to_string();
|
||||
|
||||
if trimmed.is_empty() {
|
||||
return vec![canonical];
|
||||
}
|
||||
|
||||
let mut candidates = Vec::new();
|
||||
let mut push_unique = |candidate: String| {
|
||||
if !candidates.contains(&candidate) {
|
||||
candidates.push(candidate);
|
||||
}
|
||||
};
|
||||
|
||||
push_unique(format!("{canonical}/{trimmed}"));
|
||||
push_unique(format!("/{trimmed}/.well-known/oauth-authorization-server"));
|
||||
push_unique(canonical);
|
||||
|
||||
candidates
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn discovery_paths_prefer_rfc_8414_resource_path() {
|
||||
assert_eq!(
|
||||
discovery_paths("/mcp"),
|
||||
vec![
|
||||
"/.well-known/oauth-authorization-server/mcp".to_string(),
|
||||
"/mcp/.well-known/oauth-authorization-server".to_string(),
|
||||
"/.well-known/oauth-authorization-server".to_string(),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discovery_paths_deduplicate_root_path() {
|
||||
assert_eq!(
|
||||
discovery_paths("/"),
|
||||
vec!["/.well-known/oauth-authorization-server".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_scopes_trims_empties_and_deduplicates() {
|
||||
assert_eq!(
|
||||
normalize_scopes(Some(vec![
|
||||
"read".to_string(),
|
||||
" write ".to_string(),
|
||||
"".to_string(),
|
||||
"read".to_string(),
|
||||
])),
|
||||
Some(vec!["read".to_string(), "write".to_string()])
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,7 @@ use anyhow::Context;
|
||||
use anyhow::Error;
|
||||
use anyhow::Result;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
use codex_exec_server::HttpClient;
|
||||
use oauth2::AccessToken;
|
||||
use oauth2::EmptyExtraTokenFields;
|
||||
use oauth2::RefreshToken;
|
||||
@@ -45,9 +46,9 @@ use tracing::warn;
|
||||
|
||||
use codex_keyring_store::DefaultKeyringStore;
|
||||
use codex_keyring_store::KeyringStore;
|
||||
use rmcp::transport::auth::AuthorizationManager;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::mcp_oauth_http::OAuthHttpClient;
|
||||
use codex_utils_home_dir::find_codex_home;
|
||||
|
||||
const KEYRING_SERVICE: &str = "Codex MCP Credentials";
|
||||
@@ -256,8 +257,7 @@ pub(crate) struct OAuthPersistor {
|
||||
|
||||
struct OAuthPersistorInner {
|
||||
server_name: String,
|
||||
url: String,
|
||||
authorization_manager: Arc<Mutex<AuthorizationManager>>,
|
||||
oauth_http: OAuthHttpClient,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
last_credentials: Mutex<Option<StoredOAuthTokens>>,
|
||||
}
|
||||
@@ -265,107 +265,82 @@ struct OAuthPersistorInner {
|
||||
impl OAuthPersistor {
|
||||
pub(crate) fn new(
|
||||
server_name: String,
|
||||
url: String,
|
||||
authorization_manager: Arc<Mutex<AuthorizationManager>>,
|
||||
oauth_http: OAuthHttpClient,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
initial_credentials: Option<StoredOAuthTokens>,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(OAuthPersistorInner {
|
||||
server_name,
|
||||
url,
|
||||
authorization_manager,
|
||||
oauth_http,
|
||||
store_mode,
|
||||
last_credentials: Mutex::new(initial_credentials),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn with_http_client(
|
||||
server_name: String,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
default_headers: reqwest::header::HeaderMap,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
initial_credentials: Option<StoredOAuthTokens>,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
server_name,
|
||||
OAuthHttpClient::from_default_headers(http_client, default_headers),
|
||||
store_mode,
|
||||
initial_credentials,
|
||||
)
|
||||
}
|
||||
|
||||
/// Persists the latest stored credentials if they have changed.
|
||||
/// Deletes the credentials if they are no longer present.
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "AuthorizationManager async access must be serialized through its mutex"
|
||||
)]
|
||||
pub(crate) async fn persist_if_needed(&self) -> Result<()> {
|
||||
let (client_id, maybe_credentials) = {
|
||||
let manager = self.inner.authorization_manager.clone();
|
||||
let guard = manager.lock().await;
|
||||
guard.get_credentials().await
|
||||
}?;
|
||||
|
||||
match maybe_credentials {
|
||||
Some(credentials) => {
|
||||
let mut last_credentials = self.inner.last_credentials.lock().await;
|
||||
let new_token_response = WrappedOAuthTokenResponse(credentials.clone());
|
||||
let same_token = last_credentials
|
||||
.as_ref()
|
||||
.map(|prev| prev.token_response == new_token_response)
|
||||
.unwrap_or(false);
|
||||
let expires_at = if same_token {
|
||||
last_credentials.as_ref().and_then(|prev| prev.expires_at)
|
||||
} else {
|
||||
compute_expires_at_millis(&credentials)
|
||||
};
|
||||
let stored = StoredOAuthTokens {
|
||||
server_name: self.inner.server_name.clone(),
|
||||
url: self.inner.url.clone(),
|
||||
client_id,
|
||||
token_response: new_token_response,
|
||||
expires_at,
|
||||
};
|
||||
if last_credentials.as_ref() != Some(&stored) {
|
||||
save_oauth_tokens(&self.inner.server_name, &stored, self.inner.store_mode)?;
|
||||
*last_credentials = Some(stored);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
let mut last_serialized = self.inner.last_credentials.lock().await;
|
||||
if last_serialized.take().is_some()
|
||||
&& let Err(error) = delete_oauth_tokens(
|
||||
&self.inner.server_name,
|
||||
&self.inner.url,
|
||||
self.inner.store_mode,
|
||||
)
|
||||
{
|
||||
warn!(
|
||||
"failed to remove OAuth tokens for server {}: {error}",
|
||||
self.inner.server_name
|
||||
);
|
||||
}
|
||||
}
|
||||
let last_credentials = self.inner.last_credentials.lock().await;
|
||||
if let Some(stored) = last_credentials.as_ref() {
|
||||
save_oauth_tokens(&self.inner.server_name, stored, self.inner.store_mode)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "AuthorizationManager async access must be serialized through its mutex"
|
||||
)]
|
||||
pub(crate) async fn refresh_if_needed(&self) -> Result<()> {
|
||||
let expires_at = {
|
||||
let tokens = {
|
||||
let guard = self.inner.last_credentials.lock().await;
|
||||
guard.as_ref().and_then(|tokens| tokens.expires_at)
|
||||
let Some(tokens) = guard.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
if !token_needs_refresh(tokens.expires_at) {
|
||||
return Ok(());
|
||||
}
|
||||
tokens.clone()
|
||||
};
|
||||
|
||||
if !token_needs_refresh(expires_at) {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
{
|
||||
let manager = self.inner.authorization_manager.clone();
|
||||
let guard = manager.lock().await;
|
||||
guard.refresh_token().await.with_context(|| {
|
||||
if let Some(refreshed) = self
|
||||
.inner
|
||||
.oauth_http
|
||||
.refresh_token(&tokens)
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"failed to refresh OAuth tokens for server {}",
|
||||
self.inner.server_name
|
||||
)
|
||||
})?;
|
||||
})?
|
||||
{
|
||||
let mut guard = self.inner.last_credentials.lock().await;
|
||||
*guard = Some(refreshed);
|
||||
}
|
||||
|
||||
self.persist_if_needed().await
|
||||
}
|
||||
|
||||
pub(crate) async fn access_token(&self) -> Result<Option<String>> {
|
||||
self.refresh_if_needed().await?;
|
||||
let guard = self.inner.last_credentials.lock().await;
|
||||
Ok(guard
|
||||
.as_ref()
|
||||
.map(|tokens| tokens.token_response.0.access_token().secret().to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
const FALLBACK_FILENAME: &str = ".credentials.json";
|
||||
|
||||
@@ -7,9 +7,9 @@ use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use anyhow::bail;
|
||||
use reqwest::ClientBuilder;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::ReqwestHttpClient;
|
||||
use reqwest::Url;
|
||||
use rmcp::transport::auth::OAuthState;
|
||||
use tiny_http::Response;
|
||||
use tiny_http::Server;
|
||||
use tokio::sync::oneshot;
|
||||
@@ -18,10 +18,10 @@ use urlencoding::decode;
|
||||
|
||||
use crate::StoredOAuthTokens;
|
||||
use crate::WrappedOAuthTokenResponse;
|
||||
use crate::mcp_oauth_http::OAuthAuthorizationSession;
|
||||
use crate::mcp_oauth_http::OAuthHttpClient;
|
||||
use crate::oauth::compute_expires_at_millis;
|
||||
use crate::save_oauth_tokens;
|
||||
use crate::utils::apply_default_headers;
|
||||
use crate::utils::build_default_headers;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
|
||||
struct OauthHeaders {
|
||||
@@ -80,6 +80,34 @@ pub async fn perform_oauth_login(
|
||||
oauth_resource: Option<&str>,
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
) -> Result<()> {
|
||||
perform_oauth_login_with_client(
|
||||
server_name,
|
||||
server_url,
|
||||
store_mode,
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
scopes,
|
||||
oauth_resource,
|
||||
callback_port,
|
||||
callback_url,
|
||||
Arc::new(ReqwestHttpClient),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn perform_oauth_login_with_client(
|
||||
server_name: &str,
|
||||
server_url: &str,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
scopes: &[String],
|
||||
oauth_resource: Option<&str>,
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<()> {
|
||||
perform_oauth_login_with_browser_output(
|
||||
server_name,
|
||||
@@ -92,6 +120,7 @@ pub async fn perform_oauth_login(
|
||||
callback_port,
|
||||
callback_url,
|
||||
/*emit_browser_url*/ true,
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -107,6 +136,34 @@ pub async fn perform_oauth_login_silent(
|
||||
oauth_resource: Option<&str>,
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
) -> Result<()> {
|
||||
perform_oauth_login_silent_with_client(
|
||||
server_name,
|
||||
server_url,
|
||||
store_mode,
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
scopes,
|
||||
oauth_resource,
|
||||
callback_port,
|
||||
callback_url,
|
||||
Arc::new(ReqwestHttpClient),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn perform_oauth_login_silent_with_client(
|
||||
server_name: &str,
|
||||
server_url: &str,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
scopes: &[String],
|
||||
oauth_resource: Option<&str>,
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<()> {
|
||||
perform_oauth_login_with_browser_output(
|
||||
server_name,
|
||||
@@ -119,6 +176,7 @@ pub async fn perform_oauth_login_silent(
|
||||
callback_port,
|
||||
callback_url,
|
||||
/*emit_browser_url*/ false,
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -135,6 +193,7 @@ async fn perform_oauth_login_with_browser_output(
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
emit_browser_url: bool,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<()> {
|
||||
let headers = OauthHeaders {
|
||||
http_headers,
|
||||
@@ -151,6 +210,7 @@ async fn perform_oauth_login_with_browser_output(
|
||||
callback_port,
|
||||
callback_url,
|
||||
/*timeout_secs*/ None,
|
||||
http_client,
|
||||
)
|
||||
.await?
|
||||
.finish(emit_browser_url)
|
||||
@@ -169,6 +229,36 @@ pub async fn perform_oauth_login_return_url(
|
||||
timeout_secs: Option<i64>,
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
) -> Result<OauthLoginHandle> {
|
||||
perform_oauth_login_return_url_with_client(
|
||||
server_name,
|
||||
server_url,
|
||||
store_mode,
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
scopes,
|
||||
oauth_resource,
|
||||
timeout_secs,
|
||||
callback_port,
|
||||
callback_url,
|
||||
Arc::new(ReqwestHttpClient),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn perform_oauth_login_return_url_with_client(
|
||||
server_name: &str,
|
||||
server_url: &str,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
scopes: &[String],
|
||||
oauth_resource: Option<&str>,
|
||||
timeout_secs: Option<i64>,
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<OauthLoginHandle> {
|
||||
let headers = OauthHeaders {
|
||||
http_headers,
|
||||
@@ -185,6 +275,7 @@ pub async fn perform_oauth_login_return_url(
|
||||
callback_port,
|
||||
callback_url,
|
||||
timeout_secs,
|
||||
http_client,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -329,7 +420,8 @@ impl OauthLoginHandle {
|
||||
|
||||
struct OauthLoginFlow {
|
||||
auth_url: String,
|
||||
oauth_state: OAuthState,
|
||||
oauth_http: OAuthHttpClient,
|
||||
oauth_session: Option<OAuthAuthorizationSession>,
|
||||
rx: oneshot::Receiver<CallbackResult>,
|
||||
guard: CallbackServerGuard,
|
||||
server_name: String,
|
||||
@@ -412,6 +504,7 @@ impl OauthLoginFlow {
|
||||
callback_port: Option<u16>,
|
||||
callback_url: Option<&str>,
|
||||
timeout_secs: Option<i64>,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<Self> {
|
||||
const DEFAULT_OAUTH_TIMEOUT_SECS: i64 = 300;
|
||||
|
||||
@@ -437,25 +530,19 @@ impl OauthLoginFlow {
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
} = headers;
|
||||
let default_headers = build_default_headers(http_headers, env_http_headers)?;
|
||||
let http_client = apply_default_headers(ClientBuilder::new(), &default_headers).build()?;
|
||||
|
||||
let mut oauth_state = OAuthState::new(server_url, Some(http_client)).await?;
|
||||
let scope_refs: Vec<&str> = scopes.iter().map(String::as_str).collect();
|
||||
oauth_state
|
||||
.start_authorization(&scope_refs, &redirect_uri, Some("Codex"))
|
||||
let oauth_http = OAuthHttpClient::new(http_client, http_headers, env_http_headers)?;
|
||||
let oauth_session = oauth_http
|
||||
.start_authorization(server_url, scopes, &redirect_uri, "Codex")
|
||||
.await?;
|
||||
let auth_url = append_query_param(
|
||||
&oauth_state.get_authorization_url().await?,
|
||||
"resource",
|
||||
oauth_resource,
|
||||
);
|
||||
let auth_url =
|
||||
append_query_param(&oauth_session.authorization_url, "resource", oauth_resource);
|
||||
let timeout_secs = timeout_secs.unwrap_or(DEFAULT_OAUTH_TIMEOUT_SECS).max(1);
|
||||
let timeout = Duration::from_secs(timeout_secs as u64);
|
||||
|
||||
Ok(Self {
|
||||
auth_url,
|
||||
oauth_state,
|
||||
oauth_http,
|
||||
oauth_session: Some(oauth_session),
|
||||
rx,
|
||||
guard,
|
||||
server_name: server_name.to_string(),
|
||||
@@ -503,19 +590,16 @@ impl OauthLoginFlow {
|
||||
CallbackResult::Error(error) => return Err(anyhow!(error)),
|
||||
};
|
||||
|
||||
self.oauth_state
|
||||
.handle_callback(&code, &csrf_state)
|
||||
let oauth_session = self
|
||||
.oauth_session
|
||||
.take()
|
||||
.ok_or_else(|| anyhow!("OAuth login flow was already completed"))?;
|
||||
let (client_id, credentials) = self
|
||||
.oauth_http
|
||||
.exchange_code(oauth_session, &code, &csrf_state)
|
||||
.await
|
||||
.context("failed to handle OAuth callback")?;
|
||||
|
||||
let (client_id, credentials_opt) = self
|
||||
.oauth_state
|
||||
.get_credentials()
|
||||
.await
|
||||
.context("failed to retrieve OAuth credentials")?;
|
||||
let credentials = credentials_opt
|
||||
.ok_or_else(|| anyhow!("OAuth provider did not return credentials"))?;
|
||||
|
||||
let expires_at = compute_expires_at_millis(&credentials);
|
||||
let stored = StoredOAuthTokens {
|
||||
server_name: self.server_name.clone(),
|
||||
|
||||
@@ -12,7 +12,6 @@ use std::time::Instant;
|
||||
use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use codex_api::SharedAuthProvider;
|
||||
use codex_client::build_reqwest_client_with_custom_ca;
|
||||
use codex_config::types::McpServerEnvVar;
|
||||
use codex_exec_server::HttpClient;
|
||||
use futures::FutureExt;
|
||||
@@ -45,9 +44,6 @@ use rmcp::service::RoleClient;
|
||||
use rmcp::service::RunningService;
|
||||
use rmcp::service::{self};
|
||||
use rmcp::transport::StreamableHttpClientTransport;
|
||||
use rmcp::transport::auth::AuthClient;
|
||||
use rmcp::transport::auth::AuthError;
|
||||
use rmcp::transport::auth::OAuthState;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpError;
|
||||
use serde::Deserialize;
|
||||
@@ -68,7 +64,6 @@ use crate::oauth::StoredOAuthTokens;
|
||||
use crate::stdio_server_launcher::StdioServerCommand;
|
||||
use crate::stdio_server_launcher::StdioServerLauncher;
|
||||
use crate::stdio_server_launcher::StdioServerTransport;
|
||||
use crate::utils::apply_default_headers;
|
||||
use crate::utils::build_default_headers;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
|
||||
@@ -80,7 +75,7 @@ enum PendingTransport {
|
||||
transport: StreamableHttpClientTransport<StreamableHttpClientAdapter>,
|
||||
},
|
||||
StreamableHttpWithOAuth {
|
||||
transport: StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
|
||||
transport: StreamableHttpClientTransport<StreamableHttpClientAdapter>,
|
||||
oauth_persistor: OAuthPersistor,
|
||||
},
|
||||
}
|
||||
@@ -708,11 +703,7 @@ impl RmcpClient {
|
||||
oauth_persistor,
|
||||
})
|
||||
}
|
||||
Err(err)
|
||||
if err.downcast_ref::<AuthError>().is_some_and(|auth_err| {
|
||||
matches!(auth_err, AuthError::NoAuthorizationSupport)
|
||||
}) =>
|
||||
{
|
||||
Err(err) if err.to_string().contains("No authorization support") => {
|
||||
let access_token = initial_tokens
|
||||
.token_response
|
||||
.0
|
||||
@@ -730,6 +721,7 @@ impl RmcpClient {
|
||||
Arc::clone(http_client),
|
||||
default_headers,
|
||||
/*auth_provider*/ None,
|
||||
/*oauth_persistor*/ None,
|
||||
),
|
||||
http_config,
|
||||
);
|
||||
@@ -749,6 +741,7 @@ impl RmcpClient {
|
||||
Arc::clone(http_client),
|
||||
default_headers,
|
||||
auth_provider.clone(),
|
||||
/*oauth_persistor*/ None,
|
||||
),
|
||||
http_config,
|
||||
);
|
||||
@@ -943,52 +936,27 @@ async fn create_oauth_transport_and_runtime(
|
||||
default_headers: HeaderMap,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
) -> Result<(
|
||||
StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
|
||||
StreamableHttpClientTransport<StreamableHttpClientAdapter>,
|
||||
OAuthPersistor,
|
||||
)> {
|
||||
let builder = apply_default_headers(reqwest::Client::builder(), &default_headers);
|
||||
let oauth_metadata_client = build_reqwest_client_with_custom_ca(builder)?;
|
||||
// TODO(aibrahim): teach OAuth bootstrap and refresh to use the same
|
||||
// shared HTTP client abstraction instead of always creating the local
|
||||
// reqwest metadata client here.
|
||||
let mut oauth_state =
|
||||
OAuthState::new(url.to_string(), Some(oauth_metadata_client.clone())).await?;
|
||||
|
||||
oauth_state
|
||||
.set_credentials(
|
||||
&initial_tokens.client_id,
|
||||
initial_tokens.token_response.0.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let manager = match oauth_state {
|
||||
OAuthState::Authorized(manager) => manager,
|
||||
OAuthState::Unauthorized(manager) => manager,
|
||||
OAuthState::Session(_) | OAuthState::AuthorizedHttpClient(_) => {
|
||||
return Err(anyhow!("unexpected OAuth state during client setup"));
|
||||
}
|
||||
};
|
||||
|
||||
let auth_client = AuthClient::new(
|
||||
StreamableHttpClientAdapter::new(http_client, default_headers, /*auth_provider*/ None),
|
||||
manager,
|
||||
);
|
||||
let auth_manager = auth_client.auth_manager.clone();
|
||||
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
auth_client,
|
||||
StreamableHttpClientTransportConfig::with_uri(url.to_string()),
|
||||
);
|
||||
|
||||
let runtime = OAuthPersistor::new(
|
||||
let oauth_persistor = OAuthPersistor::with_http_client(
|
||||
server_name.to_string(),
|
||||
url.to_string(),
|
||||
auth_manager,
|
||||
Arc::clone(&http_client),
|
||||
default_headers.clone(),
|
||||
credentials_store,
|
||||
Some(initial_tokens),
|
||||
);
|
||||
|
||||
Ok((transport, runtime))
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
StreamableHttpClientAdapter::new(
|
||||
http_client,
|
||||
default_headers,
|
||||
/*auth_provider*/ None,
|
||||
Some(oauth_persistor.clone()),
|
||||
),
|
||||
StreamableHttpClientTransportConfig::with_uri(url.to_string()),
|
||||
);
|
||||
Ok((transport, oauth_persistor))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use codex_config::types::McpServerEnvVar;
|
||||
use reqwest::ClientBuilder;
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest::header::HeaderName;
|
||||
use reqwest::header::HeaderValue;
|
||||
@@ -115,17 +114,6 @@ pub(crate) fn build_default_headers(
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
pub(crate) fn apply_default_headers(
|
||||
builder: ClientBuilder,
|
||||
default_headers: &HeaderMap,
|
||||
) -> ClientBuilder {
|
||||
if default_headers.is_empty() {
|
||||
builder
|
||||
} else {
|
||||
builder.default_headers(default_headers.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
pub(crate) const DEFAULT_ENV_VARS: &[&str] = &[
|
||||
"HOME",
|
||||
|
||||
@@ -0,0 +1,309 @@
|
||||
//! Regression coverage for remote-environment Streamable HTTP OAuth.
|
||||
//!
|
||||
//! The OAuth issuer in this test uses an unresolvable hostname. If any
|
||||
//! discovery, registration, or token-exchange request bypasses the injected
|
||||
//! `HttpClient`, the test fails before the callback can complete.
|
||||
|
||||
mod streamable_http_test_support;
|
||||
|
||||
use std::ffi::OsString;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::MutexGuard;
|
||||
use std::sync::OnceLock;
|
||||
use std::sync::PoisonError;
|
||||
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
use codex_exec_server::ExecServerError;
|
||||
use codex_exec_server::HttpClient;
|
||||
use codex_exec_server::HttpHeader;
|
||||
use codex_exec_server::HttpRequestParams;
|
||||
use codex_exec_server::HttpRequestResponse;
|
||||
use codex_exec_server::HttpResponseBodyStream;
|
||||
use codex_rmcp_client::StoredOAuthTokens;
|
||||
use codex_rmcp_client::WrappedOAuthTokenResponse;
|
||||
use codex_rmcp_client::perform_oauth_login_return_url_with_client;
|
||||
use codex_rmcp_client::save_oauth_tokens;
|
||||
use futures::FutureExt as _;
|
||||
use futures::future::BoxFuture;
|
||||
use oauth2::AccessToken;
|
||||
use oauth2::EmptyExtraTokenFields;
|
||||
use oauth2::RefreshToken;
|
||||
use oauth2::basic::BasicTokenType;
|
||||
use pretty_assertions::assert_eq;
|
||||
use rmcp::transport::auth::OAuthTokenResponse;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use serial_test::serial;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use streamable_http_test_support::call_echo_tool;
|
||||
use streamable_http_test_support::create_remote_oauth_client;
|
||||
use streamable_http_test_support::expected_echo_result;
|
||||
use streamable_http_test_support::spawn_exec_server;
|
||||
use streamable_http_test_support::spawn_streamable_http_server_with_oauth_bearer;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct RecordingHttpClient {
|
||||
requests: Arc<Mutex<Vec<HttpRequestParams>>>,
|
||||
}
|
||||
|
||||
impl RecordingHttpClient {
|
||||
fn recorded_requests(&self) -> Vec<HttpRequestParams> {
|
||||
self.requests
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpClient for RecordingHttpClient {
|
||||
fn http_request(
|
||||
&self,
|
||||
params: HttpRequestParams,
|
||||
) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
|
||||
let requests = Arc::clone(&self.requests);
|
||||
async move {
|
||||
requests
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.push(params.clone());
|
||||
|
||||
let response = match (params.method.as_str(), params.url.as_str()) {
|
||||
("GET", "http://oauth.test/.well-known/oauth-authorization-server/mcp") => {
|
||||
json_response(json!({
|
||||
"authorization_endpoint": "http://oauth.test/authorize",
|
||||
"token_endpoint": "http://oauth.test/token",
|
||||
"registration_endpoint": "http://oauth.test/register",
|
||||
"scopes_supported": ["tools.read", "tools.write"],
|
||||
"response_types_supported": ["code"],
|
||||
}))?
|
||||
}
|
||||
("POST", "http://oauth.test/register") => json_response(json!({
|
||||
"client_id": "registered-client",
|
||||
}))?,
|
||||
("POST", "http://oauth.test/token") => json_response(json!({
|
||||
"access_token": "remote-access-token",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
"refresh_token": "remote-refresh-token",
|
||||
}))?,
|
||||
_ => {
|
||||
return Err(ExecServerError::HttpRequest(format!(
|
||||
"unexpected HTTP request: {} {}",
|
||||
params.method, params.url
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
|
||||
fn http_request_stream(
|
||||
&self,
|
||||
params: HttpRequestParams,
|
||||
) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
|
||||
async move {
|
||||
Err(ExecServerError::HttpRequest(format!(
|
||||
"unexpected streaming HTTP request: {} {}",
|
||||
params.method, params.url
|
||||
)))
|
||||
}
|
||||
.boxed()
|
||||
}
|
||||
}
|
||||
|
||||
fn json_response(body: Value) -> Result<HttpRequestResponse, ExecServerError> {
|
||||
let body = serde_json::to_vec(&body)
|
||||
.map_err(|err| ExecServerError::HttpRequest(format!("serialize JSON response: {err}")))?;
|
||||
Ok(HttpRequestResponse {
|
||||
status: 200,
|
||||
headers: vec![HttpHeader {
|
||||
name: "content-type".to_string(),
|
||||
value: "application/json".to_string(),
|
||||
}],
|
||||
body: body.into(),
|
||||
})
|
||||
}
|
||||
|
||||
struct TempCodexHome {
|
||||
_guard: MutexGuard<'static, ()>,
|
||||
previous: Option<OsString>,
|
||||
_dir: TempDir,
|
||||
}
|
||||
|
||||
impl TempCodexHome {
|
||||
fn new() -> anyhow::Result<Self> {
|
||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
let guard = LOCK
|
||||
.get_or_init(Mutex::default)
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
let previous = std::env::var_os("CODEX_HOME");
|
||||
let dir = TempDir::new()?;
|
||||
unsafe {
|
||||
std::env::set_var("CODEX_HOME", dir.path());
|
||||
}
|
||||
Ok(Self {
|
||||
_guard: guard,
|
||||
previous,
|
||||
_dir: dir,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TempCodexHome {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
match self.previous.as_ref() {
|
||||
Some(value) => std::env::set_var("CODEX_HOME", value),
|
||||
None => std::env::remove_var("CODEX_HOME"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn browser_callback_flow_uses_injected_http_client_for_oauth_requests() -> anyhow::Result<()>
|
||||
{
|
||||
let _codex_home = TempCodexHome::new()?;
|
||||
let http_client = RecordingHttpClient::default();
|
||||
|
||||
let handle = perform_oauth_login_return_url_with_client(
|
||||
"remote-oauth-test",
|
||||
"http://oauth.test/mcp",
|
||||
OAuthCredentialsStoreMode::File,
|
||||
/*http_headers*/ None,
|
||||
/*env_http_headers*/ None,
|
||||
&["tools.read".to_string()],
|
||||
/*oauth_resource*/ None,
|
||||
Some(10),
|
||||
/*callback_port*/ None,
|
||||
/*callback_url*/ None,
|
||||
Arc::new(http_client.clone()),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let authorization_url = reqwest::Url::parse(handle.authorization_url())?;
|
||||
let mut state = None;
|
||||
let mut redirect_uri = None;
|
||||
for (name, value) in authorization_url.query_pairs() {
|
||||
match name.as_ref() {
|
||||
"state" => state = Some(value.into_owned()),
|
||||
"redirect_uri" => redirect_uri = Some(value.into_owned()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let state = state.expect("authorization URL includes state");
|
||||
let redirect_uri = redirect_uri.expect("authorization URL includes redirect_uri");
|
||||
let callback_url = format!("{redirect_uri}?code=provider-code&state={state}");
|
||||
let callback_response = reqwest::get(callback_url).await?;
|
||||
assert_eq!(callback_response.status(), reqwest::StatusCode::OK);
|
||||
|
||||
handle.wait().await?;
|
||||
|
||||
let requests = http_client.recorded_requests();
|
||||
assert_eq!(
|
||||
requests
|
||||
.iter()
|
||||
.map(|request| (request.method.as_str(), request.url.as_str()))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
(
|
||||
"GET",
|
||||
"http://oauth.test/.well-known/oauth-authorization-server/mcp"
|
||||
),
|
||||
("POST", "http://oauth.test/register"),
|
||||
("POST", "http://oauth.test/token"),
|
||||
]
|
||||
);
|
||||
|
||||
let registration_body: Value = serde_json::from_slice(
|
||||
&requests[1]
|
||||
.body
|
||||
.as_ref()
|
||||
.expect("registration request has a body")
|
||||
.0,
|
||||
)?;
|
||||
assert_eq!(
|
||||
registration_body,
|
||||
json!({
|
||||
"client_name": "Codex",
|
||||
"redirect_uris": [redirect_uri],
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"token_endpoint_auth_method": "none",
|
||||
"response_types": ["code"],
|
||||
})
|
||||
);
|
||||
|
||||
let token_body = String::from_utf8(
|
||||
requests[2]
|
||||
.body
|
||||
.as_ref()
|
||||
.expect("token request has a body")
|
||||
.0
|
||||
.clone(),
|
||||
)?;
|
||||
assert!(token_body.contains("grant_type=authorization_code"));
|
||||
assert!(token_body.contains("code=provider-code"));
|
||||
assert!(token_body.contains("client_id=registered-client"));
|
||||
assert!(token_body.contains("code_verifier="));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stored_oauth_refreshes_and_authenticates_remote_streamable_http() -> anyhow::Result<()> {
|
||||
let _codex_home = TempCodexHome::new()?;
|
||||
let refresh_token = "remote-refresh-token";
|
||||
let refreshed_access_token = "refreshed-remote-access-token";
|
||||
let (_server, base_url) =
|
||||
spawn_streamable_http_server_with_oauth_bearer(refreshed_access_token, refresh_token)
|
||||
.await?;
|
||||
let exec_server = spawn_exec_server().await?;
|
||||
|
||||
save_oauth_tokens(
|
||||
"test-streamable-http-remote-oauth",
|
||||
&expired_oauth_tokens(
|
||||
"test-streamable-http-remote-oauth",
|
||||
&format!("{base_url}/mcp"),
|
||||
"expired-access-token",
|
||||
refresh_token,
|
||||
),
|
||||
OAuthCredentialsStoreMode::File,
|
||||
)?;
|
||||
|
||||
let client = create_remote_oauth_client(&base_url, exec_server.client.clone()).await?;
|
||||
let result = call_echo_tool(&client, "remote-oauth").await?;
|
||||
|
||||
assert_eq!(result, expected_echo_result("remote-oauth"));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn expired_oauth_tokens(
|
||||
server_name: &str,
|
||||
url: &str,
|
||||
access_token: &str,
|
||||
refresh_token: &str,
|
||||
) -> StoredOAuthTokens {
|
||||
let mut token_response = OAuthTokenResponse::new(
|
||||
AccessToken::new(access_token.to_string()),
|
||||
BasicTokenType::Bearer,
|
||||
EmptyExtraTokenFields {},
|
||||
);
|
||||
token_response.set_refresh_token(Some(RefreshToken::new(refresh_token.to_string())));
|
||||
|
||||
StoredOAuthTokens {
|
||||
server_name: server_name.to_string(),
|
||||
url: url.to_string(),
|
||||
client_id: "stored-client".to_string(),
|
||||
token_response: WrappedOAuthTokenResponse(token_response),
|
||||
expires_at: Some(0),
|
||||
}
|
||||
}
|
||||
@@ -128,10 +128,40 @@ pub(crate) async fn create_remote_client(
|
||||
base_url: &str,
|
||||
http_client: ExecServerClient,
|
||||
) -> anyhow::Result<RmcpClient> {
|
||||
let client = RmcpClient::new_streamable_http_client(
|
||||
create_remote_client_with_bearer(
|
||||
"test-streamable-http-remote",
|
||||
&format!("{base_url}/mcp"),
|
||||
base_url,
|
||||
Some("test-bearer".to_string()),
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Creates a Streamable HTTP RMCP client that authenticates using stored OAuth
|
||||
/// credentials through the remote runtime HTTP API.
|
||||
pub(crate) async fn create_remote_oauth_client(
|
||||
base_url: &str,
|
||||
http_client: ExecServerClient,
|
||||
) -> anyhow::Result<RmcpClient> {
|
||||
create_remote_client_with_bearer(
|
||||
"test-streamable-http-remote-oauth",
|
||||
base_url,
|
||||
/*bearer_token*/ None,
|
||||
http_client,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn create_remote_client_with_bearer(
|
||||
server_name: &str,
|
||||
base_url: &str,
|
||||
bearer_token: Option<String>,
|
||||
http_client: ExecServerClient,
|
||||
) -> anyhow::Result<RmcpClient> {
|
||||
let client = RmcpClient::new_streamable_http_client(
|
||||
server_name,
|
||||
&format!("{base_url}/mcp"),
|
||||
bearer_token,
|
||||
/*http_headers*/ None,
|
||||
/*env_http_headers*/ None,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
@@ -193,16 +223,38 @@ pub(crate) async fn arm_session_post_failure(
|
||||
}
|
||||
|
||||
pub(crate) async fn spawn_streamable_http_server() -> anyhow::Result<(Child, String)> {
|
||||
spawn_streamable_http_server_with_env(&[]).await
|
||||
}
|
||||
|
||||
pub(crate) async fn spawn_streamable_http_server_with_oauth_bearer(
|
||||
bearer_token: &str,
|
||||
refresh_token: &str,
|
||||
) -> anyhow::Result<(Child, String)> {
|
||||
spawn_streamable_http_server_with_env(&[
|
||||
("MCP_EXPECT_BEARER", bearer_token),
|
||||
("MCP_OAUTH_ACCESS_TOKEN", bearer_token),
|
||||
("MCP_OAUTH_REFRESH_TOKEN", refresh_token),
|
||||
])
|
||||
.await
|
||||
}
|
||||
|
||||
async fn spawn_streamable_http_server_with_env(
|
||||
env: &[(&str, &str)],
|
||||
) -> anyhow::Result<(Child, String)> {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")?;
|
||||
let port = listener.local_addr()?.port();
|
||||
drop(listener);
|
||||
|
||||
let bind_addr = format!("127.0.0.1:{port}");
|
||||
let base_url = format!("http://{bind_addr}");
|
||||
let mut child = Command::new(streamable_http_server_bin()?)
|
||||
let mut command = Command::new(streamable_http_server_bin()?);
|
||||
command
|
||||
.kill_on_drop(true)
|
||||
.env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr)
|
||||
.spawn()?;
|
||||
.env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr);
|
||||
for (name, value) in env {
|
||||
command.env(name, value);
|
||||
}
|
||||
let mut child = command.spawn()?;
|
||||
|
||||
wait_for_streamable_http_server(&mut child, &bind_addr, Duration::from_secs(5)).await?;
|
||||
Ok((child, base_url))
|
||||
|
||||
Reference in New Issue
Block a user