openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-13 11:47:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Handle paginated MCP tool discovery

Browse Source
This commit is contained in:
Daniel Steigman
2026-06-09 20:13:03 -07:00
parent 4be1a168fc
commit a91f6ac3ab
3 changed files with 317 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
codex-rs/codex-mcp/src/rmcp_client.rs
View File

@@ -8,6 +8,7 @@
use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::HashSet;
use std::env;
use std::ffi::OsString;
use std::sync::Arc;
@@ -61,6 +62,7 @@ use rmcp::model::ClientCapabilities;
use rmcp::model::ElicitationCapability;
use rmcp::model::Implementation;
use rmcp::model::InitializeRequestParams;
use rmcp::model::PaginatedRequestParams;
use rmcp::model::ProtocolVersion;
use rmcp::model::Tool as RmcpTool;
use tokio_util::sync::CancellationToken;
@@ -345,13 +347,18 @@ pub(crate) async fn list_tools_for_client_uncached(
timeout: Option<Duration>,
server_instructions: Option<&str>,
) -> Result<Vec<ToolInfo>> {
let resp = client
.list_tools_with_connector_ids(/*params*/ None, timeout)
.await?;
let tools = resp
.tools
.into_iter()
.map(|tool| {
let mut tools = Vec::new();
let mut cursor: Option<String> = None;
let mut seen_cursors = HashSet::new();
loop {
let params = cursor
.as_ref()
.map(|next| PaginatedRequestParams::default().with_cursor(Some(next.clone())));
let resp = client
.list_tools_with_connector_ids(params, timeout)
.await?;
tools.extend(resp.tools.into_iter().map(|tool| {
let mut tool_def = tool.tool;
let (connector_id, connector_name, connector_description) =
sanitize_tool_connector_metadata(
@@ -396,12 +403,28 @@ pub(crate) async fn list_tools_for_client_uncached(
connector_name,
plugin_display_names: Vec::new(),
}
})
.collect();
if server_name == CODEX_APPS_MCP_SERVER_NAME {
return Ok(filter_disallowed_codex_apps_tools(tools));
}));
match resp.next_cursor {
Some(next) => {
cursor = Some(advance_tools_list_cursor(&mut seen_cursors, next)?);
}
None => {
if server_name == CODEX_APPS_MCP_SERVER_NAME {
return Ok(filter_disallowed_codex_apps_tools(tools));
}
return Ok(tools);
}
}
}
}
fn advance_tools_list_cursor(seen_cursors: &mut HashSet<String>, next: String) -> Result<String> {
if seen_cursors.insert(next.clone()) {
Ok(next)
} else {
Err(anyhow!("tools/list returned duplicate cursor"))
}
Ok(tools)
}
fn sanitize_tool_connector_metadata(
@@ -660,6 +683,7 @@ async fn make_rmcp_client(
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
use rmcp::model::JsonObject;
use rmcp::model::Meta;
@@ -753,4 +777,24 @@ mod tests {
assert!(meta.0.contains_key(key), "{key} should be preserved");
}
}
#[test]
fn tools_list_cursor_guard_rejects_cursor_cycles() {
let mut seen_cursors = HashSet::new();
assert_eq!(
advance_tools_list_cursor(&mut seen_cursors, "page-a".to_string())
.expect("first cursor should be accepted"),
"page-a"
);
assert_eq!(
advance_tools_list_cursor(&mut seen_cursors, "page-b".to_string())
.expect("second cursor should be accepted"),
"page-b"
);
let error = advance_tools_list_cursor(&mut seen_cursors, "page-a".to_string())
.expect_err("repeated cursor should be rejected");
assert_eq!(error.to_string(), "tools/list returned duplicate cursor");
}
}

146
codex-rs/core/tests/suite/rmcp_client.rs
View File

@@ -1753,6 +1753,7 @@ async fn remote_stdio_env_var_source_does_not_copy_local_env() -> anyhow::Result
const REMOTE_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_TEST_REMOTE_EXEC_SERVER_URL";
/// OAuth metadata path served by the Streamable HTTP MCP test server.
const STREAMABLE_HTTP_METADATA_PATH: &str = "/.well-known/oauth-authorization-server/mcp";
const TOOL_LIST_SESSION_IDS_PATH: &str = "/test/state/tool-list-session-ids";
/// Streamable HTTP test server plus the process handle needed for cleanup.
struct StreamableHttpTestServer {
@@ -1766,6 +1767,12 @@ enum StreamableHttpTestServerProcess {
Remote(RemoteStreamableHttpServer),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ToolPagination {
Disabled,
Enabled,
}
/// Remote Streamable HTTP server process and copied files to remove on drop.
struct RemoteStreamableHttpServer {
container_name: String,
@@ -1872,8 +1879,12 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> {
// placement. In full CI this may be the remote environment container; locally
// it is a host process.
let expected_env_value = "propagated-env-http";
let Some(http_server) =
start_streamable_http_test_server(expected_env_value, /*expected_token*/ None).await?
let Some(http_server) = start_streamable_http_test_server(
expected_env_value,
/*expected_token*/ None,
ToolPagination::Disabled,
)
.await?
else {
return Ok(());
};
@@ -1973,6 +1984,88 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
#[serial(streamable_http_env)]
async fn streamable_http_bearer_env_var_discovers_paginated_tools() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let mock = mount_sse_once(
&server,
responses::sse(vec![
responses::ev_response_created("resp-1"),
responses::ev_assistant_message("msg-1", "ready"),
responses::ev_completed("resp-1"),
]),
)
.await;
let expected_token = "test-bearer";
let token_env_var = "CODEX_TEST_STREAMABLE_HTTP_MCP_TOKEN";
let _token_guard = EnvVarGuard::set(token_env_var, OsStr::new(expected_token));
let Some(http_server) = start_streamable_http_test_server(
"paginated-tools-env",
Some(expected_token),
ToolPagination::Enabled,
)
.await?
else {
return Ok(());
};
let server_url = http_server.url().to_string();
let session_state_server_url = server_url.clone();
let server_name = "rmcp_http_paginated";
let namespace = format!("mcp__{server_name}");
let fixture = test_codex()
.with_config(move |config| {
insert_mcp_server(
config,
server_name,
McpServerTransportConfig::StreamableHttp {
url: server_url,
bearer_token_env_var: Some(token_env_var.to_string()),
http_headers: None,
env_http_headers: None,
},
TestMcpServerOptions {
environment_id: remote_aware_environment_id(),
..Default::default()
},
);
})
.build_with_remote_env(&server)
.await?;
wait_for_mcp_server(&fixture.codex, server_name).await?;
assert_recorded_tool_list_session_continuity(&session_state_server_url).await?;
fixture
.codex
.submit(read_only_user_turn(
&fixture,
"confirm the paginated streamable HTTP tools are available",
))
.await?;
wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let request = mock.single_request();
assert!(
request.tool_by_name(&namespace, "echo").is_some(),
"first-page streamable HTTP MCP tool should be exposed"
);
assert!(
request
.tool_by_name(&namespace, "second_page_tool")
.is_some(),
"second-page streamable HTTP MCP tool should be exposed"
);
http_server.shutdown().await;
server.verify().await;
Ok(())
}
/// This test writes to a fallback credentials file in CODEX_HOME.
/// Ideally, we wouldn't need to serialize the test but it's much more cumbersome to wire CODEX_HOME through the code.
#[test]
@@ -2043,8 +2136,12 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
let expected_token = "initial-access-token";
let client_id = "test-client-id";
let refresh_token = "initial-refresh-token";
let Some(http_server) =
start_streamable_http_test_server(expected_env_value, Some(expected_token)).await?
let Some(http_server) = start_streamable_http_test_server(
expected_env_value,
Some(expected_token),
ToolPagination::Disabled,
)
.await?
else {
return Ok(());
};
@@ -2167,6 +2264,7 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
async fn start_streamable_http_test_server(
expected_env_value: &str,
expected_token: Option<&str>,
tool_pagination: ToolPagination,
) -> anyhow::Result<Option<StreamableHttpTestServer>> {
let rmcp_http_server_bin = match cargo_bin("test_streamable_http_server") {
Ok(path) => path,
@@ -2183,6 +2281,7 @@ async fn start_streamable_http_test_server(
&rmcp_http_server_bin,
expected_env_value,
expected_token,
tool_pagination,
)
.await?,
));
@@ -2202,6 +2301,9 @@ async fn start_streamable_http_test_server(
if let Some(expected_token) = expected_token {
command.env("MCP_EXPECT_BEARER", expected_token);
}
if tool_pagination == ToolPagination::Enabled {
command.env("MCP_PAGINATE_TOOLS", "1");
}
let mut child = command.spawn()?;
wait_for_local_streamable_http_server(&mut child, &server_url, Duration::from_secs(5)).await?;
@@ -2217,6 +2319,7 @@ async fn start_remote_streamable_http_test_server(
rmcp_http_server_bin: &Path,
expected_env_value: &str,
expected_token: Option<&str>,
tool_pagination: ToolPagination,
) -> anyhow::Result<StreamableHttpTestServer> {
let remote_path = copy_binary_to_remote_env(
container_name,
@@ -2242,6 +2345,9 @@ async fn start_remote_streamable_http_test_server(
sh_single_quote(expected_token)
));
}
if tool_pagination == ToolPagination::Enabled {
env_assignments.push("MCP_PAGINATE_TOOLS=1".to_string());
}
let script = format!(
"{} nohup {} > {} 2>&1 < /dev/null & echo $!",
@@ -2512,6 +2618,38 @@ fn streamable_http_metadata_url(server_url: &str) -> String {
format!("{base_url}{STREAMABLE_HTTP_METADATA_PATH}")
}
async fn assert_recorded_tool_list_session_continuity(server_url: &str) -> anyhow::Result<()> {
let base_url = server_url.strip_suffix("/mcp").unwrap_or(server_url);
let session_ids: Vec<Option<String>> = Client::builder()
.no_proxy()
.build()?
.get(format!("{base_url}{TOOL_LIST_SESSION_IDS_PATH}"))
.send()
.await?
.error_for_status()?
.json()
.await?;
ensure!(
session_ids.len() >= 2,
"expected at least two paginated tools/list calls, got {}",
session_ids.len()
);
let first = session_ids
.first()
.and_then(Option::as_deref)
.filter(|value| !value.is_empty())
.ok_or_else(|| anyhow::anyhow!("first tools/list request did not include a session id"))?;
ensure!(
session_ids
.iter()
.all(|session_id| session_id.as_deref() == Some(first)),
"paginated tools/list calls did not reuse one initialized session"
);
Ok(())
}
fn write_fallback_oauth_tokens(
home: &Path,
server_name: &str,

129
codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs
View File

@@ -8,6 +8,7 @@ use std::time::Duration;
use axum::Router;
use axum::body::Body;
use axum::body::to_bytes;
use axum::extract::Json;
use axum::extract::State;
use axum::http::HeaderMap;
@@ -21,6 +22,7 @@ use axum::http::header::HOST;
use axum::http::header::WWW_AUTHENTICATE;
use axum::middleware;
use axum::middleware::Next;
use axum::response::IntoResponse;
use axum::response::Response;
use axum::routing::get;
use axum::routing::post;
@@ -64,12 +66,24 @@ const MEMO_URI: &str = "memo://codex/example-note";
const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp test server.";
const MCP_SESSION_ID_HEADER: &str = "mcp-session-id";
const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure";
const TOOL_LIST_SESSION_IDS_PATH: &str = "/test/state/tool-list-session-ids";
#[derive(Clone, Default)]
struct TestServerState {
session_failure: SessionFailureState,
tool_list_sessions: ToolListSessionState,
}
#[derive(Clone, Default)]
struct SessionFailureState {
armed_failure: Arc<Mutex<Option<ArmedFailure>>>,
}
#[derive(Clone, Default)]
struct ToolListSessionState {
session_ids: Arc<Mutex<Vec<Option<String>>>>,
}
#[derive(Clone, Debug)]
struct ArmedFailure {
status: StatusCode,
@@ -97,7 +111,7 @@ struct EchoArgs {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let bind_addr = parse_bind_addr()?;
let session_failure_state = SessionFailureState::default();
let test_state = TestServerState::default();
const MAX_BIND_RETRIES: u32 = 20;
const BIND_RETRY_DELAY: Duration = Duration::from_millis(50);
@@ -129,6 +143,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
SESSION_POST_FAILURE_CONTROL_PATH,
post(arm_session_post_failure),
)
.route(TOOL_LIST_SESSION_IDS_PATH, get(tool_list_session_ids))
.route(
"/.well-known/oauth-authorization-server/mcp",
get({
@@ -162,10 +177,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
),
)
.layer(middleware::from_fn_with_state(
session_failure_state.clone(),
test_state.clone(),
fail_session_post_when_armed,
))
.with_state(session_failure_state);
.layer(middleware::from_fn_with_state(
test_state.clone(),
record_tool_list_session_ids,
))
.with_state(test_state);
let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") {
let expected = Arc::new(format!("Bearer {token}"));
@@ -192,11 +211,31 @@ impl ServerHandler for TestToolServer {
fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
request: Option<PaginatedRequestParams>,
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> impl std::future::Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
let tools = self.tools.clone();
async move {
if std::env::var_os("MCP_PAGINATE_TOOLS").is_some() {
let cursor = request.as_ref().and_then(|params| params.cursor.as_deref());
return match cursor {
None => Ok(ListToolsResult {
tools: tools.iter().take(1).cloned().collect(),
next_cursor: Some("page-2".to_string()),
meta: None,
}),
Some("page-2") => Ok(ListToolsResult {
tools: tools.iter().skip(1).cloned().collect(),
next_cursor: None,
meta: None,
}),
Some(cursor) => Err(McpError::invalid_params(
format!("unknown tool cursor: {cursor}"),
None,
)),
};
}
Ok(ListToolsResult {
tools: (*tools).clone(),
next_cursor: None,
@@ -294,7 +333,7 @@ impl ServerHandler for TestToolServer {
impl TestToolServer {
fn new() -> Self {
let tools = vec![Self::echo_tool()];
let tools = vec![Self::echo_tool(), Self::second_page_tool()];
let resources = vec![Self::memo_resource()];
let resource_templates = vec![Self::memo_template()];
Self {
@@ -343,6 +382,24 @@ impl TestToolServer {
tool
}
fn second_page_tool() -> Tool {
#[expect(clippy::expect_used)]
let schema: JsonObject = serde_json::from_value(json!({
"type": "object",
"properties": {},
"additionalProperties": false
}))
.expect("second-page tool schema should deserialize");
let mut tool = Tool::new(
Cow::Borrowed("second_page_tool"),
Cow::Borrowed("Tool that only appears on the second tools/list page."),
Arc::new(schema),
);
tool.annotations = Some(ToolAnnotations::new().read_only(true));
tool
}
fn memo_resource() -> Resource {
let raw = RawResource {
uri: MEMO_URI.to_string(),
@@ -389,7 +446,8 @@ async fn require_bearer(
request: Request<Body>,
next: Next,
) -> Result<Response, StatusCode> {
if request.uri().path().contains("/.well-known/") {
if request.uri().path().contains("/.well-known/") || request.uri().path().starts_with("/test/")
{
return Ok(next.run(request).await);
}
if request
@@ -404,7 +462,7 @@ async fn require_bearer(
}
async fn arm_session_post_failure(
State(state): State<SessionFailureState>,
State(state): State<TestServerState>,
Json(request): Json<ArmSessionPostFailureRequest>,
) -> Result<StatusCode, StatusCode> {
let status = StatusCode::from_u16(request.status).map_err(|_| StatusCode::BAD_REQUEST)?;
@@ -422,12 +480,16 @@ async fn arm_session_post_failure(
www_authenticate_headers,
})
};
*state.armed_failure.lock().await = armed_failure;
*state.session_failure.armed_failure.lock().await = armed_failure;
Ok(StatusCode::NO_CONTENT)
}
async fn tool_list_session_ids(State(state): State<TestServerState>) -> Json<Vec<Option<String>>> {
Json(state.tool_list_sessions.session_ids.lock().await.clone())
}
async fn fail_session_post_when_armed(
State(state): State<SessionFailureState>,
State(state): State<TestServerState>,
request: Request<Body>,
next: Next,
) -> Response {
@@ -439,7 +501,7 @@ async fn fail_session_post_when_armed(
}
{
let mut armed_failure = state.armed_failure.lock().await;
let mut armed_failure = state.session_failure.armed_failure.lock().await;
if let Some(failure) = armed_failure.as_mut()
&& failure.remaining > 0
{
@@ -464,3 +526,50 @@ async fn fail_session_post_when_armed(
next.run(request).await
}
async fn record_tool_list_session_ids(
State(state): State<TestServerState>,
request: Request<Body>,
next: Next,
) -> Response {
if request.uri().path() != "/mcp" || request.method() != Method::POST {
return next.run(request).await;
}
let session_id = request
.headers()
.get(MCP_SESSION_ID_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let (parts, body) = request.into_parts();
let body_bytes = match to_bytes(body, 1024 * 1024).await {
Ok(body) => body,
Err(error) => {
return (
StatusCode::BAD_REQUEST,
format!("failed to read request body: {error}"),
)
.into_response();
}
};
if serde_json::from_slice::<serde_json::Value>(&body_bytes)
.ok()
.and_then(|body| {
body.get("method")
.and_then(serde_json::Value::as_str)
.map(str::to_string)
})
.as_deref()
== Some("tools/list")
{
state
.tool_list_sessions
.session_ids
.lock()
.await
.push(session_id);
}
next.run(Request::from_parts(parts, Body::from(body_bytes)))
.await
}