[codex] Support protected resource OAuth discovery

codex-rs/app-server/tests/suite/v2/plugin_install.rs

@@ -1330,6 +1330,97 @@ async fn plugin_install_starts_mcp_oauth_with_formerly_disallowed_plugin_app() -
     Ok(())
 }
 
+#[tokio::test]
+async fn plugin_install_starts_mcp_oauth_through_protected_resource_metadata() -> Result<()> {
+    let resource_server = MockServer::start().await;
+    let authorization_server = MockServer::start().await;
+    let resource_metadata_url = format!("{}/oauth-resource", resource_server.uri());
+    let challenge = format!("Bearer resource_metadata=\"{resource_metadata_url}\"");
+    Mock::given(method("GET"))
+        .and(path("/mcp"))
+        .respond_with(
+            ResponseTemplate::new(401).insert_header("WWW-Authenticate", challenge.as_str()),
+        )
+        .mount(&resource_server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path("/oauth-resource"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "resource": resource_server.uri(),
+            "authorization_servers": [authorization_server.uri()],
+        })))
+        .mount(&resource_server)
+        .await;
+    Mock::given(method("GET"))
+        .and(path("/.well-known/oauth-authorization-server"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "authorization_endpoint": format!("{}/oauth/authorize", authorization_server.uri()),
+            "token_endpoint": format!("{}/oauth/token", authorization_server.uri()),
+            "registration_endpoint": format!("{}/oauth/register", authorization_server.uri()),
+            "response_types_supported": ["code"],
+            "code_challenge_methods_supported": ["S256"],
+        })))
+        .mount(&authorization_server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/register"))
+        .respond_with(ResponseTemplate::new(400))
+        .mount(&authorization_server)
+        .await;
+
+    let codex_home = TempDir::new()?;
+    std::fs::write(
+        codex_home.path().join("config.toml"),
+        "[features]\nplugins = true\n",
+    )?;
+    let repo_root = TempDir::new()?;
+    write_plugin_marketplace(
+        repo_root.path(),
+        "debug",
+        "sample-plugin",
+        "./sample-plugin",
+        /*install_policy*/ None,
+        /*auth_policy*/ None,
+    )?;
+    write_plugin_source(repo_root.path(), "sample-plugin", &[])?;
+    write_plugin_mcp_config(repo_root.path(), "sample-plugin", &resource_server.uri())?;
+    let marketplace_path =
+        AbsolutePathBuf::try_from(repo_root.path().join(".agents/plugins/marketplace.json"))?;
+
+    let mut mcp = TestAppServer::new(codex_home.path()).await?;
+    timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??;
+
+    let request_id = mcp
+        .send_plugin_install_request(PluginInstallParams {
+            marketplace_path: Some(marketplace_path),
+            remote_marketplace_name: None,
+            plugin_name: "sample-plugin".to_string(),
+        })
+        .await?;
+    let response: JSONRPCResponse = timeout(
+        DEFAULT_TIMEOUT,
+        mcp.read_stream_until_response_message(RequestId::Integer(request_id)),
+    )
+    .await??;
+    let _: PluginInstallResponse = to_response(response)?;
+    wait_for_remote_plugin_request_count(
+        &authorization_server,
+        "POST",
+        "/oauth/register",
+        /*expected_count*/ 1,
+    )
+    .await?;
+
+    let resource_metadata_requested = resource_server
+        .received_requests()
+        .await
+        .unwrap_or_default()
+        .iter()
+        .any(|request| request.url.path() == "/oauth-resource");
+    assert!(resource_metadata_requested);
+    Ok(())
+}
+
 #[tokio::test]
 async fn plugin_install_starts_mcp_oauth_for_api_key_dual_surface_plugin() -> Result<()> {
     let oauth_server = MockServer::start().await;

codex-rs/rmcp-client/src/auth_status.rs

@@ -1,15 +1,14 @@
 use std::collections::HashMap;
 use std::time::Duration;
 
-use anyhow::Error;
 use anyhow::Result;
 use codex_protocol::protocol::McpAuthStatus;
+use futures::FutureExt;
 use reqwest::Client;
-use reqwest::StatusCode;
-use reqwest::Url;
 use reqwest::header::AUTHORIZATION;
 use reqwest::header::HeaderMap;
-use serde::Deserialize;
+use rmcp::transport::AuthorizationManager;
+use rmcp::transport::auth::AuthError;
 use tracing::debug;
 
 use crate::oauth::StoredOAuthTokenStatus;
@@ -20,8 +19,6 @@ use codex_config::types::AuthKeyringBackendKind;
 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 {
@@ -89,65 +86,19 @@ async fn discover_streamable_http_oauth_with_headers(
     url: &str,
     default_headers: &HeaderMap,
 ) -> 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),
-            }));
-        }
+    let mut authorization_manager = AuthorizationManager::new(url).await?;
+    authorization_manager.with_client(client)?;
+    match authorization_manager.discover_metadata().boxed().await {
+        Ok(metadata) => Ok(Some(StreamableHttpOAuthDiscovery {
+            scopes_supported: normalize_scopes(metadata.scopes_supported),
+        })),
+        Err(AuthError::NoAuthorizationSupport) => Ok(None),
+        Err(err) => Err(err.into()),
     }
-
-    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>> {
@@ -172,37 +123,13 @@ fn normalize_scopes(scopes_supported: Option<Vec<String>>) -> Option<Vec<String>
     }
 }
 
-/// 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
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
     use axum::Json;
     use axum::Router;
+    use axum::http::StatusCode;
+    use axum::http::header::WWW_AUTHENTICATE;
     use axum::routing::get;
     use pretty_assertions::assert_eq;
     use serial_test::serial;
@@ -344,6 +271,65 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn discover_streamable_http_oauth_follows_protected_resource_metadata() {
+        let authorization_server = spawn_oauth_discovery_server(serde_json::json!({
+            "authorization_endpoint": "https://example.com/authorize",
+            "token_endpoint": "https://example.com/token",
+            "scopes_supported": ["read", " write ", "read"],
+        }))
+        .await;
+
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("listener should bind");
+        let address = listener.local_addr().expect("listener should have address");
+        let resource_metadata_url = format!("http://{address}/oauth-resource");
+        let challenge = format!("Bearer resource_metadata=\"{resource_metadata_url}\"");
+        let authorization_server_url = authorization_server.url.clone();
+        let app = Router::new()
+            .route(
+                "/mcp",
+                get(move || {
+                    let challenge = challenge.clone();
+                    async move { (StatusCode::UNAUTHORIZED, [(WWW_AUTHENTICATE, challenge)]) }
+                }),
+            )
+            .route(
+                "/oauth-resource",
+                get(move || {
+                    let authorization_server_url = authorization_server_url.clone();
+                    async move {
+                        Json(serde_json::json!({
+                            "resource": format!("http://{address}/mcp"),
+                            "authorization_servers": [authorization_server_url],
+                        }))
+                    }
+                }),
+            );
+        let handle = tokio::spawn(async move {
+            axum::serve(listener, app).await.expect("server should run");
+        });
+        let resource_server = TestServer {
+            url: format!("http://{address}/mcp"),
+            handle,
+        };
+
+        let discovery = discover_streamable_http_oauth(
+            &resource_server.url,
+            /*http_headers*/ None,
+            /*env_http_headers*/ None,
+        )
+        .await
+        .expect("discovery should succeed")
+        .expect("oauth support should be detected");
+
+        assert_eq!(
+            discovery.scopes_supported,
+            Some(vec!["read".to_string(), "write".to_string()])
+        );
+    }
+
     #[tokio::test]
     async fn discover_streamable_http_oauth_ignores_empty_scopes() {
         let server = spawn_oauth_discovery_server(serde_json::json!({