mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
Surface MCP reconnect signals when expired OAuth tokens cannot refresh (#43947)
## Why Refresh failures could leave expired MCP OAuth credentials reporting ordinary errors instead of signaling that authentication is required, preventing tool calls from offering a reconnect signal. ## What changed - Classify provider failures and timeouts as `AuthorizationRequired` when the access token has expired. Keep proactive refresh failures as ordinary errors while the token remains valid. - Reread stored credentials after a failed refresh and adopt a valid login completed during the request, enforcing the existing issuer binding for refresh tokens. Preserve stored credentials for later retries. - Convert authentication-required tool-call errors into a reconnect message with `mcp/www_authenticate` metadata, without exposing provider or transport details. ## Testing Add coverage for startup and runtime OAuth recovery, proactive refresh failures, concurrent login adoption, issuer mismatches, credential preservation, and successful retries. Verify failed local refreshes do not send tool calls and server-rejected calls are not replayed. GitOrigin-RevId: 9cf64b0452685d2085df5dc0bcecee16ffe42e47
This commit is contained in:
committed by
copyberry
parent
7c098d8741
commit
5ac0b8768d
@@ -5,6 +5,7 @@ use std::time::SystemTime;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Error;
|
||||
use anyhow::Result;
|
||||
use codex_keyring_store::DefaultKeyringStore;
|
||||
use codex_keyring_store::KeyringStore;
|
||||
@@ -24,6 +25,7 @@ use super::OAuthPersistorInner;
|
||||
use super::StoredOAuthTokens;
|
||||
use super::WrappedOAuthTokenResponse;
|
||||
use super::compute_expires_at_millis;
|
||||
use super::expires_in_from_timestamp;
|
||||
use super::refresh_lock::RefreshCredentialLock;
|
||||
use super::token_needs_refresh;
|
||||
use super::validate_refresh_token_issuer;
|
||||
@@ -155,8 +157,7 @@ impl OAuthPersistor {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Preserve RMCP's `AuthorizationRequired` marker only for credentials known to be
|
||||
// unrefreshable. Network and provider failures below remain ordinary errors.
|
||||
// Without a refresh token, authorization is required before contacting the provider.
|
||||
if !latest.has_refresh_token() {
|
||||
return Err(AuthError::AuthorizationRequired).with_context(|| {
|
||||
format!(
|
||||
@@ -192,8 +193,8 @@ impl OAuthPersistor {
|
||||
refreshed_tokens(token_response, &latest, &self.inner)
|
||||
}
|
||||
Ok(Err(error @ AuthError::TokenRefreshRejected(_))) => {
|
||||
// RMCP 3 distinguishes definitive refresh-token rejection from transient
|
||||
// provider failures. Only a rejected token requires a fresh authorization.
|
||||
// Definitive rejection requires authorization even if the access token has not
|
||||
// expired yet. Other refresh failures below check actual access-token expiry.
|
||||
warn!(
|
||||
error = %error,
|
||||
"MCP OAuth refresh token was rejected; reauthorization required"
|
||||
@@ -210,22 +211,26 @@ impl OAuthPersistor {
|
||||
error = %error,
|
||||
"MCP OAuth provider refresh failed"
|
||||
);
|
||||
return Err(error).with_context(|| {
|
||||
format!(
|
||||
"failed to refresh OAuth tokens for server {}",
|
||||
self.inner.server_name
|
||||
)
|
||||
});
|
||||
let error = Error::new(error).context(format!(
|
||||
"failed to refresh OAuth tokens for server {}",
|
||||
self.inner.server_name
|
||||
));
|
||||
return self
|
||||
.recover_after_failed_refresh(keyring_store, &mut guard, &latest, error)
|
||||
.await;
|
||||
}
|
||||
Err(_) => {
|
||||
warn!(
|
||||
timeout_ms = refresh_request_timeout.as_millis(),
|
||||
"MCP OAuth provider refresh timed out; the outcome is unknown and a later serialized retry is permitted"
|
||||
);
|
||||
anyhow::bail!(
|
||||
let error = anyhow::anyhow!(
|
||||
"timed out after {refresh_request_timeout:?} refreshing OAuth tokens for server {}",
|
||||
self.inner.server_name
|
||||
);
|
||||
return self
|
||||
.recover_after_failed_refresh(keyring_store, &mut guard, &latest, error)
|
||||
.await;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -267,6 +272,63 @@ impl OAuthPersistor {
|
||||
debug!("persisted refreshed MCP OAuth credentials and completed the transaction");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn recover_after_failed_refresh<K: KeyringStore + Clone + 'static>(
|
||||
&self,
|
||||
keyring_store: &K,
|
||||
manager: &mut AuthorizationManager,
|
||||
previous: &StoredOAuthTokens,
|
||||
error: Error,
|
||||
) -> Result<()> {
|
||||
// A failed proactive refresh does not require a new login while the access token is
|
||||
// still valid. Use actual expiry, not the 30-second refresh buffer.
|
||||
if !token_has_expired(previous.expires_at) {
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
// Browser login can finish while the provider request is pending. Reread the pinned
|
||||
// authority before prompting, and never delete or overwrite a replacement credential.
|
||||
let replacement = self.inner.credential_store.load(
|
||||
keyring_store,
|
||||
&self.inner.server_name,
|
||||
&self.inner.url,
|
||||
)?;
|
||||
if let Some(replacement) = replacement
|
||||
&& !token_has_expired(replacement.expires_at)
|
||||
&& !replacement.client_id.trim().is_empty()
|
||||
&& !replacement
|
||||
.token_response
|
||||
.0
|
||||
.access_token()
|
||||
.secret()
|
||||
.trim()
|
||||
.is_empty()
|
||||
{
|
||||
// Match the existing pre-refresh adoption rule: refresh credentials must remain
|
||||
// bound to the issuer already validated for this authorization manager.
|
||||
if replacement.has_refresh_token()
|
||||
&& (replacement.bound_issuer().is_none()
|
||||
|| replacement.bound_issuer() != previous.bound_issuer())
|
||||
{
|
||||
return Err(AuthError::AuthorizationRequired).context(
|
||||
"replacement MCP OAuth refresh credentials do not match the validated issuer",
|
||||
);
|
||||
}
|
||||
debug!("adopting new MCP OAuth credentials after a failed refresh");
|
||||
install_tokens_in_manager(manager, &replacement).await?;
|
||||
*self.inner.last_credentials.lock().await = Some(replacement);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
warn!("MCP OAuth access token is expired and refresh failed; reauthorization required");
|
||||
// Keep AuthorizationRequired as the source so both startup classification and runtime
|
||||
// tool-call recovery recognize it; retain the original failure as diagnostic context.
|
||||
Err(Error::new(AuthError::AuthorizationRequired).context(error))
|
||||
}
|
||||
}
|
||||
|
||||
fn token_has_expired(expires_at: Option<u64>) -> bool {
|
||||
expires_at.is_some_and(|expires_at| expires_in_from_timestamp(expires_at).is_none())
|
||||
}
|
||||
|
||||
/// Installs tokens without resolving metadata again, so callers can pin the validated snapshot.
|
||||
|
||||
@@ -49,6 +49,7 @@ use crate::oauth::OAuthRuntime;
|
||||
use crate::oauth::ResolvedOAuthCredentialStore;
|
||||
use crate::oauth::StoredOAuthTokens;
|
||||
use crate::oauth::WrappedOAuthTokenResponse;
|
||||
use crate::oauth::compute_expires_at_millis;
|
||||
use crate::oauth::compute_store_key;
|
||||
use crate::oauth::delete_oauth_tokens;
|
||||
use crate::oauth::load_oauth_tokens_from_file;
|
||||
@@ -533,8 +534,14 @@ async fn issuerless_newer_credentials_are_not_adopted_before_refresh() -> Result
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn transient_refresh_failure_does_not_require_reauthorization() -> Result<()> {
|
||||
let (_env, server, initial) = test_context().await?;
|
||||
async fn proactive_refresh_failure_with_unexpired_token_does_not_require_reauthorization()
|
||||
-> Result<()> {
|
||||
let (_env, server, mut initial) = test_context().await?;
|
||||
initial
|
||||
.token_response
|
||||
.0
|
||||
.set_expires_in(Some(&Duration::from_secs(/*secs*/ 30)));
|
||||
initial.expires_at = compute_expires_at_millis(&initial.token_response.0);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
@@ -638,6 +645,139 @@ async fn assert_caller_cancellation(mode: crate::McpOAuthRefreshMode) -> Result<
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn failed_refresh_adopts_login_completed_during_request() -> Result<()> {
|
||||
for expires_in in [Some(Duration::from_secs(/*secs*/ 3600)), None] {
|
||||
let (_env, server, initial) = test_context().await?;
|
||||
let mut replacement = initial.clone();
|
||||
replacement.client_id = "new-login-client".to_string();
|
||||
replacement
|
||||
.token_response
|
||||
.0
|
||||
.set_access_token(AccessToken::new("new-login-access-token".to_string()));
|
||||
replacement
|
||||
.token_response
|
||||
.0
|
||||
.set_expires_in(expires_in.as_ref());
|
||||
replacement.expires_at = compute_expires_at_millis(&replacement.token_response.0);
|
||||
save_oauth_tokens_to_file(&initial)?;
|
||||
let persistor = persistor_for(&initial).await?;
|
||||
let provider_replacement = replacement.clone();
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
.respond_with(move |_request: &wiremock::Request| {
|
||||
// A browser login writes independently of the refresh transaction lock.
|
||||
save_oauth_tokens_to_file(&provider_replacement).expect("complete new login");
|
||||
ResponseTemplate::new(200).set_body_string("not an OAuth token response")
|
||||
})
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
persistor.refresh_if_needed().await?;
|
||||
let live = persistor.stored_credentials().await.expect("adopted login");
|
||||
let mut expected_live = replacement.clone();
|
||||
expected_live.token_response.0.set_expires_in(None);
|
||||
assert_eq!(live, expected_live);
|
||||
persistor.persist_if_needed().await?;
|
||||
let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?
|
||||
.expect("new login must remain stored");
|
||||
assert_tokens_match_without_expiry(&stored, &replacement);
|
||||
server.verify().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn failed_refresh_does_not_adopt_unbound_replacement_credentials() -> Result<()> {
|
||||
for issuer in [None, Some("https://different-issuer.example.test")] {
|
||||
let (_env, server, initial) = test_context().await?;
|
||||
let mut replacement = initial.clone();
|
||||
replacement.issuer = issuer.map(str::to_string);
|
||||
replacement.expires_at = Some(u64::MAX);
|
||||
replacement
|
||||
.token_response
|
||||
.0
|
||||
.set_access_token(AccessToken::new("replacement-access-token".to_string()));
|
||||
save_oauth_tokens_to_file(&initial)?;
|
||||
let persistor = persistor_for(&initial).await?;
|
||||
let provider_replacement = replacement.clone();
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
.respond_with(move |_request: &wiremock::Request| {
|
||||
save_oauth_tokens_to_file(&provider_replacement).expect("complete new login");
|
||||
ResponseTemplate::new(200).set_body_string("invalid JSON")
|
||||
})
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let error = persistor
|
||||
.refresh_if_needed()
|
||||
.await
|
||||
.expect_err("an unbound replacement must not enter the authorization manager");
|
||||
assert!(is_authentication_required_error(&error));
|
||||
let live = persistor
|
||||
.stored_credentials()
|
||||
.await
|
||||
.expect("original credentials");
|
||||
let mut expected_live = initial.clone();
|
||||
expected_live.token_response.0.set_expires_in(None);
|
||||
assert_eq!(live, expected_live);
|
||||
let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?
|
||||
.expect("the replacement must not be deleted or overwritten");
|
||||
assert_tokens_match_without_expiry(&stored, &replacement);
|
||||
server.verify().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn expired_refresh_failure_preserves_credentials_for_a_later_retry() -> Result<()> {
|
||||
let (_env, server, initial) = test_context().await?;
|
||||
save_oauth_tokens_to_file(&initial)?;
|
||||
let persistor = persistor_for(&initial).await?;
|
||||
let failure = Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string("invalid JSON"))
|
||||
.expect(1)
|
||||
.mount_as_scoped(&server)
|
||||
.await;
|
||||
|
||||
let error = persistor
|
||||
.refresh_if_needed()
|
||||
.await
|
||||
.expect_err("refresh failed");
|
||||
assert!(is_authentication_required_error(&error));
|
||||
let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?
|
||||
.expect("failed refresh must preserve stored credentials");
|
||||
assert_tokens_match_without_expiry(&stored, &initial);
|
||||
drop(failure);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.and(body_string_contains("refresh_token=refresh-token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "retry-access-token",
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
persistor.refresh_if_needed().await?;
|
||||
let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)?
|
||||
.expect("successful retry must persist new credentials");
|
||||
assert_eq!(
|
||||
stored.token_response.0.access_token().secret(),
|
||||
"retry-access-token"
|
||||
);
|
||||
server.verify().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> Result<()> {
|
||||
let (_env, server, initial) = test_context().await?;
|
||||
@@ -653,6 +793,7 @@ async fn provider_timeout_releases_lock_and_preserves_durable_credentials() -> R
|
||||
.await
|
||||
.expect_err("provider request should reach its explicit timeout");
|
||||
assert!(error.to_string().contains("timed out after 50ms"));
|
||||
assert!(is_authentication_required_error(&error));
|
||||
|
||||
let _lock = tokio::time::timeout(
|
||||
Duration::from_millis(/*millis*/ 100),
|
||||
|
||||
@@ -92,6 +92,7 @@ use crate::oauth::validate_refresh_token_issuer;
|
||||
use crate::oauth_http_client::OAuthHttpClientAdapter;
|
||||
use crate::oauth_refresh_mode::McpOAuthRefreshMode;
|
||||
use crate::protocol_mode::McpProtocolMode;
|
||||
use crate::startup_error::is_authentication_required_error;
|
||||
use crate::stdio_server_launcher::StdioServerCommand;
|
||||
use crate::stdio_server_launcher::StdioServerLauncher;
|
||||
use crate::stdio_server_launcher::StdioServerProcessHandle;
|
||||
@@ -798,7 +799,27 @@ impl RmcpClient {
|
||||
meta: Option<serde_json::Value>,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<CallToolResult> {
|
||||
self.refresh_oauth_if_needed().await?;
|
||||
let authentication_required_result = |error| {
|
||||
if !is_authentication_required_error(&error) {
|
||||
return Err(error);
|
||||
}
|
||||
// Local expiry and server rejection use the same reconnect signal without
|
||||
// exposing token-endpoint or transport details in the tool result.
|
||||
let mut result = CallToolResult::error(vec![ContentBlock::text(
|
||||
"MCP authentication required. Reconnect to continue using this server.",
|
||||
)]);
|
||||
result.meta = Some(
|
||||
serde_json::Map::from_iter([(
|
||||
"mcp/www_authenticate".to_string(),
|
||||
Value::String("Bearer error=\"invalid_token\"".to_string()),
|
||||
)])
|
||||
.into(),
|
||||
);
|
||||
Ok(result)
|
||||
};
|
||||
if let Err(error) = self.refresh_oauth_if_needed().await {
|
||||
return authentication_required_result(error);
|
||||
}
|
||||
let arguments = match arguments {
|
||||
Some(Value::Object(map)) => Some(map),
|
||||
Some(other) => {
|
||||
@@ -863,14 +884,14 @@ impl RmcpClient {
|
||||
let Some(ClientOperationError::Service(ServiceError::TransportSend(transport))) =
|
||||
error.downcast_ref()
|
||||
else {
|
||||
return Err(error);
|
||||
return authentication_required_result(error);
|
||||
};
|
||||
let Some(StreamableHttpError::AuthRequired(challenge)) =
|
||||
transport
|
||||
.error
|
||||
.downcast_ref::<StreamableHttpError<StreamableHttpClientAdapterError>>()
|
||||
else {
|
||||
return Err(error);
|
||||
return authentication_required_result(error);
|
||||
};
|
||||
// The transport has already handled automatic refresh. Preserve the challenge
|
||||
// for interactive login without replaying the rejected tool call.
|
||||
|
||||
@@ -0,0 +1,387 @@
|
||||
//! Exercises OAuth recovery through a real MCP client, with isolated file credentials.
|
||||
|
||||
mod streamable_http_test_support;
|
||||
|
||||
use std::time::Duration;
|
||||
use std::time::SystemTime;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use codex_config::types::AuthKeyringBackendKind;
|
||||
use codex_config::types::OAuthCredentialsStoreMode;
|
||||
use codex_exec_server::Environment;
|
||||
use codex_rmcp_client::RmcpClient;
|
||||
use codex_rmcp_client::StoredOAuthTokens;
|
||||
use codex_rmcp_client::WrappedOAuthTokenResponse;
|
||||
use codex_rmcp_client::is_authentication_required_error;
|
||||
use codex_rmcp_client::save_oauth_tokens;
|
||||
use codex_rmcp_client::stored_oauth_credentials;
|
||||
use oauth2::AccessToken;
|
||||
use oauth2::RefreshToken;
|
||||
use oauth2::basic::BasicTokenType;
|
||||
use pretty_assertions::assert_eq;
|
||||
use rmcp::model::CallToolResult;
|
||||
use rmcp::model::ContentBlock;
|
||||
use rmcp::transport::auth::OAuthTokenResponse;
|
||||
use rmcp::transport::auth::VendorExtraTokenFields;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tokio::process::Command;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::Request;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::body_string_contains;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
use streamable_http_test_support::initialize_client;
|
||||
|
||||
const SERVER_NAME: &str = "runtime-oauth";
|
||||
const SCENARIO_ENV: &str = "MCP_TEST_RUNTIME_OAUTH_SCENARIO";
|
||||
|
||||
#[tokio::test]
|
||||
async fn runtime_oauth_recovery() -> anyhow::Result<()> {
|
||||
run_scenarios(&[
|
||||
"no-refresh",
|
||||
"rejected",
|
||||
"malformed",
|
||||
"transient",
|
||||
"early-malformed",
|
||||
"early-transient",
|
||||
"refreshable",
|
||||
"unexpired",
|
||||
"unknown-expiry",
|
||||
"revoked",
|
||||
])
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn startup_oauth_recovery() -> anyhow::Result<()> {
|
||||
run_scenarios(&[
|
||||
"startup-malformed",
|
||||
"startup-transient",
|
||||
"startup-refreshable",
|
||||
])
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_scenarios(scenarios: &[&str]) -> anyhow::Result<()> {
|
||||
for scenario in scenarios {
|
||||
let home = TempDir::new()?;
|
||||
let output = Command::new(std::env::current_exe()?)
|
||||
.args(["runtime_oauth_child", "--exact", "--ignored", "--nocapture"])
|
||||
.env("CODEX_HOME", home.path())
|
||||
.env(SCENARIO_ENV, scenario)
|
||||
.output()
|
||||
.await?;
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{scenario}: {}\n{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "spawned by OAuth recovery tests with an isolated CODEX_HOME"]
|
||||
async fn runtime_oauth_child() -> anyhow::Result<()> {
|
||||
let scenario = std::env::var(SCENARIO_ENV)?;
|
||||
let mcp = MockServer::start().await;
|
||||
let authorization = MockServer::start().await;
|
||||
let server_url = format!("{}/mcp", mcp.uri());
|
||||
let startup = scenario.starts_with("startup-");
|
||||
let revoked = scenario == "revoked";
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/.well-known/oauth-authorization-server/mcp"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"issuer": server_url,
|
||||
"authorization_endpoint": format!("{}/authorize", authorization.uri()),
|
||||
"token_endpoint": format!("{}/token", authorization.uri()),
|
||||
})))
|
||||
.mount(&mcp)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/mcp"))
|
||||
.respond_with(move |request: &Request| {
|
||||
let body: Value = request.body_json().expect("JSON-RPC request");
|
||||
match body["method"].as_str() {
|
||||
Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0", "id": body["id"],
|
||||
"result": {
|
||||
"protocolVersion": body["params"]["protocolVersion"],
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "runtime-oauth", "version": "1"},
|
||||
},
|
||||
})),
|
||||
Some("notifications/initialized") => ResponseTemplate::new(202),
|
||||
Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0", "id": body["id"], "result": {"tools": []},
|
||||
})),
|
||||
Some("tools/call") if revoked => ResponseTemplate::new(401)
|
||||
.insert_header("www-authenticate", "Bearer error=\"invalid_token\""),
|
||||
Some("tools/call") => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0", "id": body["id"],
|
||||
"result": {"content": [{"type": "text", "text": "tool ran"}]},
|
||||
})),
|
||||
_ => ResponseTemplate::new(400),
|
||||
}
|
||||
})
|
||||
.mount(&mcp)
|
||||
.await;
|
||||
// A zero lifetime expires after the handshake; 30 seconds exercises the early-refresh
|
||||
// window while the token remains valid. Neither case needs sleeps or clock changes.
|
||||
let expires_in = match scenario.as_str() {
|
||||
"unexpired" | "startup-refreshable" => Some(3600),
|
||||
"early-malformed" | "early-transient" => Some(30),
|
||||
"unknown-expiry" => None,
|
||||
_ => Some(0),
|
||||
};
|
||||
let bootstrap_response = match scenario.as_str() {
|
||||
"startup-malformed" => ResponseTemplate::new(200).set_body_raw("{", "application/json"),
|
||||
"startup-transient" => {
|
||||
ResponseTemplate::new(503).set_body_json(json!({"error": "temporarily_unavailable"}))
|
||||
}
|
||||
_ => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"access_token": "initial-token", "token_type": "Bearer",
|
||||
"refresh_token": "refresh-token", "expires_in": expires_in,
|
||||
})),
|
||||
};
|
||||
let bootstrap = Mock::given(method("POST"))
|
||||
.and(path("/token"))
|
||||
.respond_with(bootstrap_response)
|
||||
.mount_as_scoped(&authorization)
|
||||
.await;
|
||||
let mut response = OAuthTokenResponse::new(
|
||||
AccessToken::new("old-token".to_string()),
|
||||
BasicTokenType::Bearer,
|
||||
VendorExtraTokenFields::default(),
|
||||
);
|
||||
if !revoked {
|
||||
response.set_refresh_token(Some(RefreshToken::new("refresh-token".to_string())));
|
||||
}
|
||||
let tokens = StoredOAuthTokens {
|
||||
server_name: SERVER_NAME.to_string(),
|
||||
url: server_url.clone(),
|
||||
issuer: Some(server_url.clone()),
|
||||
client_id: "test-client".to_string(),
|
||||
token_response: WrappedOAuthTokenResponse(response),
|
||||
expires_at: Some(if revoked {
|
||||
u64::try_from(SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis())? + 3_600_000
|
||||
} else {
|
||||
0
|
||||
}),
|
||||
};
|
||||
save_oauth_tokens(
|
||||
SERVER_NAME,
|
||||
&tokens,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::default(),
|
||||
)
|
||||
.await?;
|
||||
let client = RmcpClient::new_streamable_http_client(
|
||||
SERVER_NAME,
|
||||
&server_url,
|
||||
/*bearer_token*/ None,
|
||||
/*http_headers*/ None,
|
||||
/*env_http_headers*/ None,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::default(),
|
||||
Environment::default_for_tests().get_http_client(),
|
||||
/*auth_provider*/ None,
|
||||
)
|
||||
.await?;
|
||||
let initialization = initialize_client(&client).await;
|
||||
if startup {
|
||||
if scenario == "startup-refreshable" {
|
||||
initialization?;
|
||||
let tools = client
|
||||
.list_tools(/*params*/ None, Some(Duration::from_secs(5)))
|
||||
.await?;
|
||||
assert!(tools.tools.is_empty());
|
||||
} else {
|
||||
let error = initialization.expect_err("expired credentials must require reconnect");
|
||||
assert!(is_authentication_required_error(&error), "{error:#}");
|
||||
}
|
||||
let requests = mcp.received_requests().await.expect("request recording");
|
||||
let rpc_requests = requests
|
||||
.iter()
|
||||
.filter(|request| request.method.as_str() == "POST" && request.url.path() == "/mcp")
|
||||
.collect::<Vec<_>>();
|
||||
let methods = rpc_requests
|
||||
.iter()
|
||||
.map(|request| {
|
||||
request.body_json::<Value>().expect("JSON-RPC request")["method"].clone()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
methods,
|
||||
if scenario == "startup-refreshable" {
|
||||
vec![
|
||||
json!("initialize"),
|
||||
json!("notifications/initialized"),
|
||||
json!("tools/list"),
|
||||
]
|
||||
} else {
|
||||
Vec::<Value>::new()
|
||||
}
|
||||
);
|
||||
for request in rpc_requests {
|
||||
assert_eq!(
|
||||
request
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("Bearer initial-token")
|
||||
);
|
||||
}
|
||||
let after = stored_oauth_credentials(
|
||||
SERVER_NAME,
|
||||
&server_url,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::default(),
|
||||
)?
|
||||
.expect("credentials are not deleted");
|
||||
if scenario != "startup-refreshable" {
|
||||
assert_eq!(after, tokens);
|
||||
}
|
||||
assert_eq!(
|
||||
authorization
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("request recording")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
initialization?;
|
||||
drop(bootstrap);
|
||||
let mut before = stored_oauth_credentials(
|
||||
SERVER_NAME,
|
||||
&server_url,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::default(),
|
||||
)?
|
||||
.expect("saved credentials");
|
||||
if scenario == "no-refresh" {
|
||||
before.token_response.0.set_refresh_token(None);
|
||||
save_oauth_tokens(
|
||||
SERVER_NAME,
|
||||
&before,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::default(),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let refresh_response = match scenario.as_str() {
|
||||
"refreshable" => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"access_token": "renewed-token", "token_type": "Bearer", "expires_in": 3600,
|
||||
})),
|
||||
"malformed" | "early-malformed" => {
|
||||
ResponseTemplate::new(200).set_body_raw("{", "application/json")
|
||||
}
|
||||
"transient" | "early-transient" => {
|
||||
ResponseTemplate::new(503).set_body_json(json!({"error": "temporarily_unavailable"}))
|
||||
}
|
||||
_ => ResponseTemplate::new(400).set_body_json(json!({
|
||||
"error": "invalid_grant", "error_description": "private provider details",
|
||||
})),
|
||||
};
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
.respond_with(refresh_response)
|
||||
.expect(
|
||||
if matches!(
|
||||
scenario.as_str(),
|
||||
"no-refresh" | "unexpired" | "unknown-expiry" | "revoked"
|
||||
) {
|
||||
0
|
||||
} else {
|
||||
1
|
||||
},
|
||||
)
|
||||
.mount(&authorization)
|
||||
.await;
|
||||
|
||||
let result = client
|
||||
.call_tool(
|
||||
"echo".to_string(),
|
||||
/*arguments*/ None,
|
||||
/*meta*/ None,
|
||||
Some(Duration::from_secs(5)),
|
||||
)
|
||||
.await;
|
||||
match scenario.as_str() {
|
||||
"no-refresh" | "rejected" | "malformed" | "transient" => {
|
||||
let mut expected = CallToolResult::error(vec![ContentBlock::text(
|
||||
"MCP authentication required. Reconnect to continue using this server.",
|
||||
)]);
|
||||
expected.meta = Some(
|
||||
serde_json::Map::from_iter([(
|
||||
"mcp/www_authenticate".to_string(),
|
||||
json!("Bearer error=\"invalid_token\""),
|
||||
)])
|
||||
.into(),
|
||||
);
|
||||
assert_eq!(result?, expected);
|
||||
}
|
||||
"revoked" => {
|
||||
let mut expected =
|
||||
CallToolResult::error(vec![ContentBlock::text("Authentication required")]);
|
||||
expected.meta = Some(
|
||||
serde_json::Map::from_iter([(
|
||||
"mcp/www_authenticate".to_string(),
|
||||
json!(["Bearer error=\"invalid_token\""]),
|
||||
)])
|
||||
.into(),
|
||||
);
|
||||
assert_eq!(result?, expected);
|
||||
}
|
||||
"early-malformed" | "early-transient" => {
|
||||
let error = result.expect_err("early refresh failures must remain ordinary errors");
|
||||
assert!(!is_authentication_required_error(&error));
|
||||
}
|
||||
"refreshable" | "unexpired" | "unknown-expiry" => {
|
||||
let result = result?;
|
||||
assert_eq!(result.content, vec![ContentBlock::text("tool ran")]);
|
||||
assert_eq!(result.meta, None);
|
||||
}
|
||||
_ => unreachable!("scenario selected by parent test"),
|
||||
}
|
||||
let tool_calls = mcp
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("request recording")
|
||||
.iter()
|
||||
.filter(|request| {
|
||||
request
|
||||
.body_json::<Value>()
|
||||
.ok()
|
||||
.is_some_and(|body| body["method"] == "tools/call")
|
||||
})
|
||||
.count();
|
||||
assert_eq!(
|
||||
tool_calls,
|
||||
usize::from(matches!(
|
||||
scenario.as_str(),
|
||||
"refreshable" | "unexpired" | "unknown-expiry" | "revoked"
|
||||
))
|
||||
);
|
||||
let after = stored_oauth_credentials(
|
||||
SERVER_NAME,
|
||||
&server_url,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
AuthKeyringBackendKind::default(),
|
||||
)?
|
||||
.expect("credentials are not deleted");
|
||||
if scenario != "refreshable" {
|
||||
assert_eq!(before, after);
|
||||
}
|
||||
authorization.verify().await;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user