mirror of
https://github.com/openai/codex.git
synced 2026-11-10 21:00:47 +00:00
test(rmcp): cover live refresh token rotation
This commit is contained in:
@@ -37,10 +37,13 @@ const EXPIRED_ACCESS_TOKEN: &str = "expired-access-token";
|
||||
const REFRESH_TOKEN: &str = "valid-refresh-token";
|
||||
const REFRESHED_ACCESS_TOKEN: &str = "refreshed-access-token";
|
||||
const SECOND_REFRESHED_ACCESS_TOKEN: &str = "second-refreshed-access-token";
|
||||
const THIRD_REFRESHED_ACCESS_TOKEN: &str = "third-refreshed-access-token";
|
||||
const REPLACEMENT_REFRESH_TOKEN: &str = "replacement-refresh-token";
|
||||
const CHILD_SERVER_URL_ENV: &str = "MCP_TEST_OAUTH_STARTUP_SERVER_URL";
|
||||
const OMITTED_REFRESH_TOKEN_CHILD_SERVER_URL_ENV: &str =
|
||||
"MCP_TEST_OAUTH_OMITTED_REFRESH_TOKEN_SERVER_URL";
|
||||
const LIVE_REFRESH_TOKEN_CHILD_SERVER_URL_ENV: &str =
|
||||
"MCP_TEST_OAUTH_LIVE_REFRESH_TOKEN_SERVER_URL";
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn refreshes_expired_persisted_token_before_initialize() -> anyhow::Result<()> {
|
||||
@@ -383,3 +386,198 @@ async fn oauth_preserved_refresh_token_child() -> anyhow::Result<()> {
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn omitted_refresh_token_preserves_running_manager_state() -> anyhow::Result<()> {
|
||||
fn respond_to_mcp(request: &Request) -> ResponseTemplate {
|
||||
let body: Value = request.body_json().expect("valid JSON-RPC request");
|
||||
match body.get("method").and_then(Value::as_str) {
|
||||
Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body.get("id").cloned().unwrap_or(Value::Null),
|
||||
"result": {
|
||||
"protocolVersion": body
|
||||
.pointer("/params/protocolVersion")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!("2025-06-18")),
|
||||
"capabilities": {},
|
||||
"serverInfo": {
|
||||
"name": "oauth-live-refresh-token-test",
|
||||
"version": "0.0.0-test",
|
||||
},
|
||||
},
|
||||
})),
|
||||
Some("notifications/initialized") => ResponseTemplate::new(202),
|
||||
Some("tools/list") => ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body.get("id").cloned().unwrap_or(Value::Null),
|
||||
"result": {
|
||||
"tools": [],
|
||||
},
|
||||
})),
|
||||
method => ResponseTemplate::new(400)
|
||||
.set_body_string(format!("unexpected JSON-RPC method: {method:?}")),
|
||||
}
|
||||
}
|
||||
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/.well-known/oauth-authorization-server/mcp"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"authorization_endpoint": format!("{}/oauth/authorize", server.uri()),
|
||||
"token_endpoint": format!("{}/oauth/token", server.uri()),
|
||||
"scopes_supported": [""],
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let refresh_count = Arc::new(AtomicUsize::new(0));
|
||||
let refresh_count_for_responder = Arc::clone(&refresh_count);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.and(body_string_contains("grant_type=refresh_token"))
|
||||
.respond_with(move |request: &Request| {
|
||||
let body = String::from_utf8_lossy(&request.body);
|
||||
match refresh_count_for_responder.fetch_add(1, Ordering::SeqCst) {
|
||||
0 if body.contains(&format!("refresh_token={REFRESH_TOKEN}")) => {
|
||||
ResponseTemplate::new(200).set_body_json(json!({
|
||||
"access_token": REFRESHED_ACCESS_TOKEN,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 0,
|
||||
}))
|
||||
}
|
||||
1 if body.contains(&format!("refresh_token={REFRESH_TOKEN}")) => {
|
||||
ResponseTemplate::new(200).set_body_json(json!({
|
||||
"access_token": SECOND_REFRESHED_ACCESS_TOKEN,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 0,
|
||||
"refresh_token": REPLACEMENT_REFRESH_TOKEN,
|
||||
}))
|
||||
}
|
||||
2 if body.contains(&format!("refresh_token={REPLACEMENT_REFRESH_TOKEN}")) => {
|
||||
ResponseTemplate::new(200).set_body_json(json!({
|
||||
"access_token": THIRD_REFRESHED_ACCESS_TOKEN,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 7200,
|
||||
}))
|
||||
}
|
||||
request_count => ResponseTemplate::new(400).set_body_string(format!(
|
||||
"unexpected refresh request {request_count}: {body}"
|
||||
)),
|
||||
}
|
||||
})
|
||||
.expect(3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
for access_token in [
|
||||
REFRESHED_ACCESS_TOKEN,
|
||||
SECOND_REFRESHED_ACCESS_TOKEN,
|
||||
THIRD_REFRESHED_ACCESS_TOKEN,
|
||||
] {
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/mcp"))
|
||||
.and(header("authorization", format!("Bearer {access_token}")))
|
||||
.respond_with(respond_to_mcp)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
|
||||
let codex_home = TempDir::new()?;
|
||||
let server_url = format!("{}/mcp", server.uri());
|
||||
let status = Command::new(std::env::current_exe()?)
|
||||
.args([
|
||||
"oauth_live_refresh_token_child",
|
||||
"--exact",
|
||||
"--ignored",
|
||||
"--nocapture",
|
||||
])
|
||||
.env("CODEX_HOME", codex_home.path())
|
||||
.env(LIVE_REFRESH_TOKEN_CHILD_SERVER_URL_ENV, server_url)
|
||||
.status()
|
||||
.await?;
|
||||
assert!(
|
||||
status.success(),
|
||||
"OAuth live refresh token child failed: {status}"
|
||||
);
|
||||
assert_eq!(refresh_count.load(Ordering::SeqCst), 3);
|
||||
server.verify().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
#[ignore = "spawned by omitted_refresh_token_preserves_running_manager_state"]
|
||||
async fn oauth_live_refresh_token_child() -> anyhow::Result<()> {
|
||||
let server_url = std::env::var(LIVE_REFRESH_TOKEN_CHILD_SERVER_URL_ENV)?;
|
||||
let codex_home = std::env::var("CODEX_HOME")?;
|
||||
|
||||
let mut response = OAuthTokenResponse::new(
|
||||
AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()),
|
||||
BasicTokenType::Bearer,
|
||||
VendorExtraTokenFields::default(),
|
||||
);
|
||||
response.set_refresh_token(Some(RefreshToken::new(REFRESH_TOKEN.to_string())));
|
||||
response.set_expires_in(Some(&Duration::from_secs(7200)));
|
||||
let tokens = StoredOAuthTokens {
|
||||
server_name: SERVER_NAME.to_string(),
|
||||
url: server_url.clone(),
|
||||
client_id: "test-client-id".to_string(),
|
||||
token_response: WrappedOAuthTokenResponse(response),
|
||||
expires_at: Some(0),
|
||||
};
|
||||
save_oauth_tokens(SERVER_NAME, &tokens, OAuthCredentialsStoreMode::File)?;
|
||||
|
||||
let client = RmcpClient::new_streamable_http_client(
|
||||
SERVER_NAME,
|
||||
&server_url,
|
||||
/*bearer_token*/ None,
|
||||
/*http_headers*/ None,
|
||||
/*env_http_headers*/ None,
|
||||
OAuthCredentialsStoreMode::File,
|
||||
Environment::default_for_tests().get_http_client(),
|
||||
/*auth_provider*/ None,
|
||||
)
|
||||
.await?;
|
||||
initialize_client(&client).await?;
|
||||
|
||||
let credentials_path = std::path::Path::new(&codex_home).join(".credentials.json");
|
||||
let persisted_after_replacement: Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&credentials_path)?)?;
|
||||
let persisted_entry = persisted_after_replacement
|
||||
.as_object()
|
||||
.and_then(|entries| entries.values().next())
|
||||
.expect("one persisted OAuth credential entry");
|
||||
assert_eq!(
|
||||
persisted_entry.get("access_token"),
|
||||
Some(&json!(SECOND_REFRESHED_ACCESS_TOKEN))
|
||||
);
|
||||
assert_eq!(
|
||||
persisted_entry.get("refresh_token"),
|
||||
Some(&json!(REPLACEMENT_REFRESH_TOKEN))
|
||||
);
|
||||
|
||||
let tools = client
|
||||
.list_tools(/*params*/ None, Some(Duration::from_secs(5)))
|
||||
.await?;
|
||||
assert!(tools.tools.is_empty());
|
||||
|
||||
let persisted_after_third_refresh: Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(credentials_path)?)?;
|
||||
let persisted_entry = persisted_after_third_refresh
|
||||
.as_object()
|
||||
.and_then(|entries| entries.values().next())
|
||||
.expect("one persisted OAuth credential entry");
|
||||
assert_eq!(
|
||||
persisted_entry.get("access_token"),
|
||||
Some(&json!(THIRD_REFRESHED_ACCESS_TOKEN))
|
||||
);
|
||||
assert_eq!(
|
||||
persisted_entry.get("refresh_token"),
|
||||
Some(&json!(REPLACEMENT_REFRESH_TOKEN))
|
||||
);
|
||||
|
||||
client.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user