fix: refresh expired MCP OAuth tokens

Repro:

- Set up an HTTP MCP server with OAuth expiring tokens like Datadog
- Use it for a bit until the token expires
- Start a new codex session

Expected results:
- The token gets transparently refreshed and the MCP server continues to
work

Actual results:
- You start getting 401s and need to do `codex mcp logout` and log back 
  in.
This commit is contained in:
Javier Soto
2025-11-04 10:25:04 -08:00
parent cb6584de46
commit bc90844a63

View File

@@ -9,6 +9,7 @@ use std::time::Duration;
use anyhow::Result;
use anyhow::anyhow;
use futures::FutureExt;
use oauth2::TokenResponse;
use mcp_types::CallToolRequestParams;
use mcp_types::CallToolResult;
use mcp_types::InitializeRequestParams;
@@ -31,6 +32,7 @@ use rmcp::service::RunningService;
use rmcp::service::{self};
use rmcp::transport::StreamableHttpClientTransport;
use rmcp::transport::auth::AuthClient;
use rmcp::transport::auth::OAuthTokenResponse;
use rmcp::transport::auth::OAuthState;
use rmcp::transport::child_process::TokioChildProcess;
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
@@ -55,6 +57,8 @@ use crate::utils::convert_to_rmcp;
use crate::utils::create_env_for_mcp_server;
use crate::utils::run_with_timeout;
const REFRESH_SKEW_SECS: u64 = 60;
enum PendingTransport {
ChildProcess(TokioChildProcess),
StreamableHttp {
@@ -397,6 +401,13 @@ async fn create_oauth_transport_and_runtime(
let auth_client = AuthClient::new(http_client, manager);
let auth_manager = auth_client.auth_manager.clone();
// If the stored token is expired or about to expire, refresh before the handshake.
if should_refresh_initial_token(&initial_tokens.token_response.0) {
if let Err(err) = auth_manager.lock().await.refresh_token().await {
warn!("failed to refresh OAuth token before handshake: {err}");
}
}
let transport = StreamableHttpClientTransport::with_client(
auth_client,
StreamableHttpClientTransportConfig::with_uri(url.to_string()),
@@ -412,3 +423,10 @@ async fn create_oauth_transport_and_runtime(
Ok((transport, runtime))
}
fn should_refresh_initial_token(token: &OAuthTokenResponse) -> bool {
match token.expires_in() {
Some(duration) => duration.as_secs() <= REFRESH_SKEW_SECS,
None => token.refresh_token().is_some(),
}
}