diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 78078720e7..6714d60f8b 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -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, server_instructions: Option<&str>, ) -> Result> { - 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 = 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, next: String) -> Result { + 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"); + } } diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 36ea7d4b2c..d14b95a5e1 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -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> { 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 { 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> = 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, diff --git a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs index 2384d39473..e44a102f83 100644 --- a/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_streamable_http_server.rs @@ -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>>, } +#[derive(Clone, Default)] +struct ToolListSessionState { + session_ids: Arc>>>, +} + #[derive(Clone, Debug)] struct ArmedFailure { status: StatusCode, @@ -97,7 +111,7 @@ struct EchoArgs { #[tokio::main] async fn main() -> Result<(), Box> { 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> { 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> { ), ) .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, + request: Option, _context: rmcp::service::RequestContext, ) -> impl std::future::Future> + 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, next: Next, ) -> Result { - 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, + State(state): State, Json(request): Json, ) -> Result { 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) -> Json>> { + Json(state.tool_list_sessions.session_ids.lock().await.clone()) +} + async fn fail_session_post_when_armed( - State(state): State, + State(state): State, request: Request, 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, + request: Request, + 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::(&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 +}