From a92dc1f1b584fb2a2b3444b3f98933a8529fa420 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 22 Apr 2026 17:41:40 -0700 Subject: [PATCH] Wire remote streamable HTTP MCP Co-authored-by: Codex --- .../codex-mcp/src/mcp_connection_manager.rs | 31 +- codex-rs/core/tests/suite/rmcp_client.rs | 513 ++--------- .../rmcp_client/streamable_http_tests.rs | 843 ++++++++++++++++++ .../src/bin/test_streamable_http_server.rs | 333 +++---- 4 files changed, 1090 insertions(+), 630 deletions(-) create mode 100644 codex-rs/core/tests/suite/rmcp_client/streamable_http_tests.rs diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager.rs b/codex-rs/codex-mcp/src/mcp_connection_manager.rs index fe8bc523ea..08d83edb47 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager.rs @@ -37,6 +37,8 @@ use codex_async_utils::OrCancelExt; use codex_config::Constrained; use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::Environment; +use codex_exec_server::HttpClient; +use codex_exec_server::ReqwestHttpClient; use codex_protocol::ToolName; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::approvals::ElicitationRequestEvent; @@ -1538,7 +1540,14 @@ async fn make_rmcp_client( } = config; let remote_environment = match experimental_environment.as_deref() { None | Some("local") => false, - Some("remote") => true, + Some("remote") => { + if !runtime_environment.environment().is_remote() { + return Err(StartupOutcomeError::from(anyhow!( + "remote MCP server `{server_name}` requires a remote environment" + ))); + } + true + } Some(environment) => { return Err(StartupOutcomeError::from(anyhow!( "unsupported experimental_environment `{environment}` for MCP server `{server_name}`" @@ -1562,14 +1571,8 @@ async fn make_rmcp_client( .collect::>() }); let launcher = if remote_environment { - let exec_environment = runtime_environment.environment(); - if !exec_environment.is_remote() { - return Err(StartupOutcomeError::from(anyhow!( - "remote MCP server `{server_name}` requires a remote executor environment" - ))); - } Arc::new(ExecutorStdioServerLauncher::new( - exec_environment.get_exec_backend(), + runtime_environment.environment().get_exec_backend(), runtime_environment.fallback_cwd(), )) } else { @@ -1589,11 +1592,11 @@ async fn make_rmcp_client( env_http_headers, bearer_token_env_var, } => { - if remote_environment && !runtime_environment.environment().is_remote() { - return Err(StartupOutcomeError::from(anyhow!( - "remote MCP server `{server_name}` requires a remote environment" - ))); - } + let http_client: Arc = if remote_environment { + runtime_environment.environment().get_http_client() + } else { + Arc::new(ReqwestHttpClient) + }; let resolved_bearer_token = match resolve_bearer_token(server_name, bearer_token_env_var.as_deref()) { Ok(token) => token, @@ -1606,7 +1609,7 @@ async fn make_rmcp_client( http_headers, env_http_headers, store_mode, - runtime_environment.environment().get_http_client(), + http_client, ) .await .map_err(StartupOutcomeError::from) diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index cb6553343f..a977247c00 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -6,7 +6,6 @@ use std::collections::HashMap; use std::ffi::OsStr; use std::ffi::OsString; use std::fs; -use std::net::TcpListener; use std::path::Path; use std::path::PathBuf; use std::process::Command as StdCommand; @@ -41,7 +40,6 @@ use codex_protocol::protocol::McpToolCallBeginEvent; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::user_input::UserInput; -use codex_utils_cargo_bin::cargo_bin; use core_test_support::assert_regex_match; use core_test_support::remote_env_env_var; use core_test_support::responses; @@ -54,14 +52,9 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_event_with_timeout; -use reqwest::Client; -use reqwest::StatusCode; use serde_json::Value; use serde_json::json; use serial_test::serial; -use tempfile::tempdir; -use tokio::process::Child; -use tokio::process::Command; use tokio::time::Instant; use tokio::time::sleep; use wiremock::MockServer; @@ -115,12 +108,9 @@ fn remote_aware_experimental_environment() -> Option { /// container and return that in-container path instead. fn remote_aware_stdio_server_bin() -> anyhow::Result { let bin = stdio_server_bin()?; - let Some(container_name) = std::env::var_os(remote_env_env_var()) else { + let Some(container_name) = remote_env_container_name()? else { return Ok(bin); }; - let container_name = container_name - .into_string() - .map_err(|value| anyhow::anyhow!("remote env container name must be utf-8: {value:?}"))?; // Keep the Docker path rewrite scoped to tests that use `build_remote_aware`. // Other MCP tests still start their stdio server from the orchestrator test @@ -131,32 +121,78 @@ fn remote_aware_stdio_server_bin() -> anyhow::Result { // path instead of the host build artifact path. // Several remote-aware MCP tests can run in parallel; give each copied // binary its own path so one test cannot replace another test's executable. + copy_binary_to_remote_env(&container_name, Path::new(&bin), "test_stdio_server") +} + +/// Returns the Docker container used by remote-aware MCP tests, when active. +fn remote_env_container_name() -> anyhow::Result> { + let Some(container_name) = std::env::var_os(remote_env_env_var()) else { + return Ok(None); + }; + Ok(Some(container_name.into_string().map_err(|value| { + anyhow::anyhow!("remote env container name must be utf-8: {value:?}") + })?)) +} + +/// Builds a collision-resistant in-container path for copied test binaries. +fn unique_remote_path(binary_name: &str) -> anyhow::Result { let unique_suffix = SystemTime::now().duration_since(UNIX_EPOCH)?.as_nanos(); - let remote_path = format!( - "/tmp/codex-remote-env/test_stdio_server-{}-{unique_suffix}", + Ok(format!( + "/tmp/codex-remote-env/{binary_name}-{}-{unique_suffix}", std::process::id() + )) +} + +/// Copies a host-built helper binary into the remote test container. +fn copy_binary_to_remote_env( + container_name: &str, + host_path: &Path, + binary_name: &str, +) -> anyhow::Result { + let remote_path = unique_remote_path(binary_name)?; + let mkdir_output = StdCommand::new("docker") + .args([ + "exec", + container_name, + "mkdir", + "-p", + "/tmp/codex-remote-env", + ]) + .output() + .context("create remote MCP test binary directory")?; + ensure!( + mkdir_output.status.success(), + "docker mkdir remote MCP test binary directory failed: stdout={} stderr={}", + String::from_utf8_lossy(&mkdir_output.stdout).trim(), + String::from_utf8_lossy(&mkdir_output.stderr).trim() ); + let container_target = format!("{container_name}:{remote_path}"); let copy_output = StdCommand::new("docker") .arg("cp") - .arg(&bin) + .arg(host_path) .arg(&container_target) .output() - .with_context(|| format!("copy {bin} to remote MCP test env"))?; + .with_context(|| { + format!( + "copy {} to remote MCP test env", + host_path.to_string_lossy() + ) + })?; ensure!( copy_output.status.success(), - "docker cp test_stdio_server failed: stdout={} stderr={}", + "docker cp {binary_name} failed: stdout={} stderr={}", String::from_utf8_lossy(©_output.stdout).trim(), String::from_utf8_lossy(©_output.stderr).trim() ); let chmod_output = StdCommand::new("docker") - .args(["exec", &container_name, "chmod", "+x", remote_path.as_str()]) + .args(["exec", container_name, "chmod", "+x", remote_path.as_str()]) .output() - .context("mark remote test_stdio_server executable")?; + .with_context(|| format!("mark remote {binary_name} executable"))?; ensure!( chmod_output.status.success(), - "docker chmod test_stdio_server failed: stdout={} stderr={}", + "docker chmod {binary_name} failed: stdout={} stderr={}", String::from_utf8_lossy(&chmod_output.stdout).trim(), String::from_utf8_lossy(&chmod_output.stderr).trim() ); @@ -164,6 +200,9 @@ fn remote_aware_stdio_server_bin() -> anyhow::Result { Ok(remote_path) } +#[path = "rmcp_client/streamable_http_tests.rs"] +mod streamable_http_tests; + async fn wait_for_mcp_tool(fixture: &TestCodex, tool_name: &str) -> anyhow::Result<()> { let tools_ready_deadline = Instant::now() + Duration::from_secs(30); loop { @@ -1807,442 +1846,6 @@ async fn remote_stdio_env_var_source_does_not_copy_local_env() -> anyhow::Result Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { - skip_if_no_network!(Ok(())); - - let server = responses::start_mock_server().await; - - let call_id = "call-456"; - let server_name = "rmcp_http"; - let namespace = format!("mcp__{server_name}__"); - - mount_sse_once( - &server, - responses::sse(vec![ - responses::ev_response_created("resp-1"), - responses::ev_function_call_with_namespace( - call_id, - &namespace, - "echo", - "{\"message\":\"ping\"}", - ), - responses::ev_completed("resp-1"), - ]), - ) - .await; - mount_sse_once( - &server, - responses::sse(vec![ - responses::ev_assistant_message( - "msg-1", - "rmcp streamable http echo tool completed successfully.", - ), - responses::ev_completed("resp-2"), - ]), - ) - .await; - - let expected_env_value = "propagated-env-http"; - let rmcp_http_server_bin = match cargo_bin("test_streamable_http_server") { - Ok(path) => path, - Err(err) => { - eprintln!("test_streamable_http_server binary not available, skipping test: {err}"); - return Ok(()); - } - }; - - let listener = TcpListener::bind("127.0.0.1:0")?; - let port = listener.local_addr()?.port(); - drop(listener); - let bind_addr = format!("127.0.0.1:{port}"); - let server_url = format!("http://{bind_addr}/mcp"); - - let mut http_server_child = Command::new(&rmcp_http_server_bin) - .kill_on_drop(true) - .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr) - .env("MCP_TEST_VALUE", expected_env_value) - .spawn()?; - - wait_for_streamable_http_server(&mut http_server_child, &bind_addr, Duration::from_secs(5)) - .await?; - - let fixture = test_codex() - .with_config(move |config| { - insert_mcp_server( - config, - server_name, - McpServerTransportConfig::StreamableHttp { - url: server_url, - bearer_token_env_var: None, - http_headers: None, - env_http_headers: None, - }, - TestMcpServerOptions::default(), - ); - }) - .build(&server) - .await?; - let session_model = fixture.session_configured.model.clone(); - - fixture - .codex - .submit(Op::UserTurn { - items: vec![UserInput::Text { - text: "call the rmcp streamable http echo tool".into(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - cwd: fixture.cwd.path().to_path_buf(), - approval_policy: AskForApproval::Never, - approvals_reviewer: None, - sandbox_policy: SandboxPolicy::new_read_only_policy(), - model: session_model, - effort: None, - summary: None, - service_tier: None, - collaboration_mode: None, - personality: None, - environments: None, - }) - .await?; - - let begin_event = wait_for_event(&fixture.codex, |ev| { - matches!(ev, EventMsg::McpToolCallBegin(_)) - }) - .await; - - let EventMsg::McpToolCallBegin(begin) = begin_event else { - unreachable!("event guard guarantees McpToolCallBegin"); - }; - assert_eq!(begin.invocation.server, server_name); - assert_eq!(begin.invocation.tool, "echo"); - - let end_event = wait_for_event(&fixture.codex, |ev| { - matches!(ev, EventMsg::McpToolCallEnd(_)) - }) - .await; - let EventMsg::McpToolCallEnd(end) = end_event else { - unreachable!("event guard guarantees McpToolCallEnd"); - }; - - let result = end - .result - .as_ref() - .expect("rmcp echo tool should return success"); - assert_eq!(result.is_error, Some(false)); - assert!( - result.content.is_empty(), - "content should default to an empty array" - ); - - let structured = result - .structured_content - .as_ref() - .expect("structured content"); - let Value::Object(map) = structured else { - panic!("structured content should be an object: {structured:?}"); - }; - let echo_value = map - .get("echo") - .and_then(Value::as_str) - .expect("echo payload present"); - assert_eq!(echo_value, "ECHOING: ping"); - let env_value = map - .get("env") - .and_then(Value::as_str) - .expect("env snapshot inserted"); - assert_eq!(env_value, expected_env_value); - - wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; - - server.verify().await; - - match http_server_child.try_wait() { - Ok(Some(_)) => {} - Ok(None) => { - let _ = http_server_child.kill().await; - } - Err(error) => { - eprintln!("failed to check streamable http server status: {error}"); - let _ = http_server_child.kill().await; - } - } - if let Err(error) = http_server_child.wait().await { - eprintln!("failed to await streamable http server shutdown: {error}"); - } - - 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] -#[serial(codex_home)] -fn streamable_http_with_oauth_round_trip() -> anyhow::Result<()> { - const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024; - - let handle = std::thread::Builder::new() - .name("streamable_http_with_oauth_round_trip".to_string()) - .stack_size(TEST_STACK_SIZE_BYTES) - .spawn(|| -> anyhow::Result<()> { - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(1) - .enable_all() - .build()?; - runtime.block_on(streamable_http_with_oauth_round_trip_impl()) - })?; - - match handle.join() { - Ok(result) => result, - Err(_) => Err(anyhow::anyhow!( - "streamable_http_with_oauth_round_trip thread panicked" - )), - } -} - -#[allow(clippy::expect_used)] -async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { - skip_if_no_network!(Ok(())); - - let server = responses::start_mock_server().await; - - let call_id = "call-789"; - let server_name = "rmcp_http_oauth"; - let tool_name = format!("mcp__{server_name}__echo"); - let namespace = format!("mcp__{server_name}__"); - - mount_sse_once( - &server, - responses::sse(vec![ - responses::ev_response_created("resp-1"), - responses::ev_function_call_with_namespace( - call_id, - &namespace, - "echo", - "{\"message\":\"ping\"}", - ), - responses::ev_completed("resp-1"), - ]), - ) - .await; - mount_sse_once( - &server, - responses::sse(vec![ - responses::ev_assistant_message( - "msg-1", - "rmcp streamable http oauth echo tool completed successfully.", - ), - responses::ev_completed("resp-2"), - ]), - ) - .await; - - let expected_env_value = "propagated-env-http-oauth"; - let expected_token = "initial-access-token"; - let client_id = "test-client-id"; - let refresh_token = "initial-refresh-token"; - let rmcp_http_server_bin = match cargo_bin("test_streamable_http_server") { - Ok(path) => path, - Err(err) => { - eprintln!("test_streamable_http_server binary not available, skipping test: {err}"); - return Ok(()); - } - }; - - let listener = TcpListener::bind("127.0.0.1:0")?; - let port = listener.local_addr()?.port(); - drop(listener); - let bind_addr = format!("127.0.0.1:{port}"); - let server_url = format!("http://{bind_addr}/mcp"); - - let mut http_server_child = Command::new(&rmcp_http_server_bin) - .kill_on_drop(true) - .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr) - .env("MCP_EXPECT_BEARER", expected_token) - .env("MCP_TEST_VALUE", expected_env_value) - .spawn()?; - - wait_for_streamable_http_server(&mut http_server_child, &bind_addr, Duration::from_secs(5)) - .await?; - - let temp_home = Arc::new(tempdir()?); - let _codex_home_guard = EnvVarGuard::set("CODEX_HOME", temp_home.path().as_os_str()); - write_fallback_oauth_tokens( - temp_home.path(), - server_name, - &server_url, - client_id, - expected_token, - refresh_token, - )?; - - let fixture = test_codex() - .with_home(temp_home.clone()) - .with_config(move |config| { - // Keep OAuth credentials isolated to this test home because Bazel - // runs the full core suite in one process. - config.mcp_oauth_credentials_store_mode = serde_json::from_value(json!("file")) - .expect("`file` should deserialize as OAuthCredentialsStoreMode"); - insert_mcp_server( - config, - server_name, - McpServerTransportConfig::StreamableHttp { - url: server_url, - bearer_token_env_var: None, - http_headers: None, - env_http_headers: None, - }, - TestMcpServerOptions::default(), - ); - }) - .build(&server) - .await?; - let session_model = fixture.session_configured.model.clone(); - - wait_for_mcp_tool(&fixture, &tool_name).await?; - - fixture - .codex - .submit(Op::UserTurn { - items: vec![UserInput::Text { - text: "call the rmcp streamable http oauth echo tool".into(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - cwd: fixture.cwd.path().to_path_buf(), - approval_policy: AskForApproval::Never, - approvals_reviewer: None, - sandbox_policy: SandboxPolicy::new_read_only_policy(), - model: session_model, - effort: None, - summary: None, - service_tier: None, - collaboration_mode: None, - personality: None, - environments: None, - }) - .await?; - - let begin_event = wait_for_event(&fixture.codex, |ev| { - matches!(ev, EventMsg::McpToolCallBegin(_)) - }) - .await; - - let EventMsg::McpToolCallBegin(begin) = begin_event else { - unreachable!("event guard guarantees McpToolCallBegin"); - }; - assert_eq!(begin.invocation.server, server_name); - assert_eq!(begin.invocation.tool, "echo"); - - let end_event = wait_for_event(&fixture.codex, |ev| { - matches!(ev, EventMsg::McpToolCallEnd(_)) - }) - .await; - let EventMsg::McpToolCallEnd(end) = end_event else { - unreachable!("event guard guarantees McpToolCallEnd"); - }; - - let result = end - .result - .as_ref() - .expect("rmcp echo tool should return success"); - assert_eq!(result.is_error, Some(false)); - assert!( - result.content.is_empty(), - "content should default to an empty array" - ); - - let structured = result - .structured_content - .as_ref() - .expect("structured content"); - let Value::Object(map) = structured else { - panic!("structured content should be an object: {structured:?}"); - }; - let echo_value = map - .get("echo") - .and_then(Value::as_str) - .expect("echo payload present"); - assert_eq!(echo_value, "ECHOING: ping"); - let env_value = map - .get("env") - .and_then(Value::as_str) - .expect("env snapshot inserted"); - assert_eq!(env_value, expected_env_value); - - wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; - - server.verify().await; - - match http_server_child.try_wait() { - Ok(Some(_)) => {} - Ok(None) => { - let _ = http_server_child.kill().await; - } - Err(error) => { - eprintln!("failed to check streamable http oauth server status: {error}"); - let _ = http_server_child.kill().await; - } - } - if let Err(error) = http_server_child.wait().await { - eprintln!("failed to await streamable http oauth server shutdown: {error}"); - } - - Ok(()) -} - -async fn wait_for_streamable_http_server( - server_child: &mut Child, - address: &str, - timeout: Duration, -) -> anyhow::Result<()> { - let deadline = Instant::now() + timeout; - let metadata_url = format!("http://{address}/.well-known/oauth-authorization-server/mcp"); - let client = Client::builder().no_proxy().build()?; - loop { - if let Some(status) = server_child.try_wait()? { - return Err(anyhow::anyhow!( - "streamable HTTP server exited early with status {status}" - )); - } - - let remaining = deadline.saturating_duration_since(Instant::now()); - - if remaining.is_zero() { - return Err(anyhow::anyhow!( - "timed out waiting for streamable HTTP server metadata at {metadata_url}: deadline reached" - )); - } - - match tokio::time::timeout(remaining, client.get(&metadata_url).send()).await { - Ok(Ok(response)) if response.status() == StatusCode::OK => return Ok(()), - Ok(Ok(response)) => { - if Instant::now() >= deadline { - return Err(anyhow::anyhow!( - "timed out waiting for streamable HTTP server metadata at {metadata_url}: HTTP {}", - response.status() - )); - } - } - Ok(Err(error)) => { - if Instant::now() >= deadline { - return Err(anyhow::anyhow!( - "timed out waiting for streamable HTTP server metadata at {metadata_url}: {error}" - )); - } - } - Err(_) => { - return Err(anyhow::anyhow!( - "timed out waiting for streamable HTTP server metadata at {metadata_url}: request timed out" - )); - } - } - - sleep(Duration::from_millis(50)).await; - } -} - fn write_fallback_oauth_tokens( home: &Path, server_name: &str, diff --git a/codex-rs/core/tests/suite/rmcp_client/streamable_http_tests.rs b/codex-rs/core/tests/suite/rmcp_client/streamable_http_tests.rs new file mode 100644 index 0000000000..a0e66726ab --- /dev/null +++ b/codex-rs/core/tests/suite/rmcp_client/streamable_http_tests.rs @@ -0,0 +1,843 @@ +//! Core integration coverage for MCP Streamable HTTP servers. +//! +//! These tests live in the core RMCP suite and verify that Codex can discover, +//! connect to, and call Streamable HTTP MCP servers in both local and +//! remote-aware placements. + +use std::net::SocketAddr; +use std::net::TcpListener; +use std::path::Path; +use std::process::Command as StdCommand; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context as _; +use anyhow::ensure; +use codex_config::types::McpServerTransportConfig; +use codex_exec_server::Environment; +use codex_exec_server::HttpRequestParams; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::user_input::UserInput; +use codex_utils_cargo_bin::cargo_bin; +use core_test_support::responses; +use core_test_support::responses::mount_sse_once; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use reqwest::Client; +use reqwest::StatusCode; +use serde_json::Value; +use serde_json::json; +use serial_test::serial; +use tempfile::tempdir; +use tokio::process::Child; +use tokio::process::Command; +use tokio::time::Instant; +use tokio::time::sleep; + +use super::EnvVarGuard; +use super::TestMcpServerOptions; +use super::copy_binary_to_remote_env; +use super::insert_mcp_server; +use super::remote_aware_experimental_environment; +use super::remote_env_container_name; +use super::wait_for_mcp_tool; +use super::write_fallback_oauth_tokens; + +/// Remote runtime websocket URL used by remote-aware MCP integration tests. +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"; + +/// Streamable HTTP test server plus the process handle needed for cleanup. +struct StreamableHttpTestServer { + server_url: String, + process: StreamableHttpTestServerProcess, +} + +/// Tracks whether the Streamable HTTP test server runs on the host or remotely. +enum StreamableHttpTestServerProcess { + Local(Child), + Remote(RemoteStreamableHttpServer), +} + +/// Remote Streamable HTTP server process and copied files to remove on drop. +struct RemoteStreamableHttpServer { + container_name: String, + pid: String, + paths_to_remove: Vec, +} + +impl Drop for RemoteStreamableHttpServer { + /// Stops the remote process and removes copied test artifacts best-effort. + fn drop(&mut self) { + self.kill(); + if self.paths_to_remove.is_empty() { + return; + } + let script = format!("rm -f {}", self.paths_to_remove.join(" ")); + let _ = StdCommand::new("docker") + .args(["exec", &self.container_name, "sh", "-lc", &script]) + .output(); + } +} + +impl RemoteStreamableHttpServer { + /// Stops the remote Streamable HTTP test server process. + fn kill(&self) { + let _ = StdCommand::new("docker") + .args(["exec", &self.container_name, "kill", &self.pid]) + .output(); + } +} + +impl StreamableHttpTestServer { + /// Returns the MCP endpoint URL that Codex should connect to. + fn url(&self) -> &str { + &self.server_url + } + + /// Stops the local or remote test server and waits for local process exit. + async fn shutdown(mut self) { + match &mut self.process { + StreamableHttpTestServerProcess::Local(child) => match child.try_wait() { + Ok(Some(_)) => {} + Ok(None) => { + let _ = child.kill().await; + } + Err(error) => { + eprintln!("failed to check streamable http server status: {error}"); + let _ = child.kill().await; + } + }, + StreamableHttpTestServerProcess::Remote(server) => { + server.kill(); + } + } + if let StreamableHttpTestServerProcess::Local(child) = &mut self.process + && let Err(error) = child.wait().await + { + eprintln!("failed to await streamable http server shutdown: {error}"); + } + } +} + +/// What this tests: Codex can discover and call a Streamable HTTP MCP tool in +/// both local and remote-aware placements, and the tool observes the expected +/// environment value from the server process that actually handled the request. +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + // Phase 1: script the model responses so Codex will call the MCP echo tool + // and then complete the turn after the tool result is returned. + let server = responses::start_mock_server().await; + + let call_id = "call-456"; + let server_name = "rmcp_http"; + let namespace = format!("mcp__{server_name}__"); + + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call_with_namespace( + call_id, + &namespace, + "echo", + "{\"message\":\"ping\"}", + ), + responses::ev_completed("resp-1"), + ]), + ) + .await; + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message( + "msg-1", + "rmcp streamable http echo tool completed successfully.", + ), + responses::ev_completed("resp-2"), + ]), + ) + .await; + + // Phase 2: start the Streamable HTTP MCP test server in the active + // placement. In full CI this may be the remote executor 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? + else { + return Ok(()); + }; + let server_url = http_server.url().to_string(); + + // Phase 3: configure Codex with the Streamable HTTP MCP server and build a + // fixture that selects remote MCP placement only when the remote test + // environment is active. + let fixture = test_codex() + .with_config(move |config| { + insert_mcp_server( + config, + server_name, + McpServerTransportConfig::StreamableHttp { + url: server_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + TestMcpServerOptions { + experimental_environment: remote_aware_experimental_environment(), + ..Default::default() + }, + ); + }) + .build_remote_aware(&server) + .await?; + let session_model = fixture.session_configured.model.clone(); + + // Phase 4: submit the user turn that should trigger the MCP tool call. + fixture + .codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "call the rmcp streamable http echo tool".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: fixture.cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + approvals_reviewer: None, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + environments: None, + }) + .await?; + + // Phase 5: assert Codex begins the expected tool invocation. + let begin_event = wait_for_event(&fixture.codex, |ev| { + matches!(ev, EventMsg::McpToolCallBegin(_)) + }) + .await; + + let EventMsg::McpToolCallBegin(begin) = begin_event else { + unreachable!("event guard guarantees McpToolCallBegin"); + }; + assert_eq!(begin.invocation.server, server_name); + assert_eq!(begin.invocation.tool, "echo"); + + // Phase 6: assert the tool result proves the server handled the request and + // propagated the expected environment value. + let end_event = wait_for_event(&fixture.codex, |ev| { + matches!(ev, EventMsg::McpToolCallEnd(_)) + }) + .await; + let EventMsg::McpToolCallEnd(end) = end_event else { + unreachable!("event guard guarantees McpToolCallEnd"); + }; + + let result = end + .result + .as_ref() + .expect("rmcp echo tool should return success"); + assert_eq!(result.is_error, Some(false)); + assert!( + result.content.is_empty(), + "content should default to an empty array" + ); + + let structured = result + .structured_content + .as_ref() + .expect("structured content"); + let Value::Object(map) = structured else { + panic!("structured content should be an object: {structured:?}"); + }; + let echo_value = map + .get("echo") + .and_then(Value::as_str) + .expect("echo payload present"); + assert_eq!(echo_value, "ECHOING: ping"); + let env_value = map + .get("env") + .and_then(Value::as_str) + .expect("env snapshot inserted"); + assert_eq!(env_value, expected_env_value); + + // Phase 7: verify the scripted model calls were consumed and clean up the + // placement-aware MCP server. + wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + server.verify().await; + + http_server.shutdown().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] +#[serial(codex_home)] +fn streamable_http_with_oauth_round_trip() -> anyhow::Result<()> { + const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024; + + let handle = std::thread::Builder::new() + .name("streamable_http_with_oauth_round_trip".to_string()) + .stack_size(TEST_STACK_SIZE_BYTES) + .spawn(|| -> anyhow::Result<()> { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build()?; + runtime.block_on(streamable_http_with_oauth_round_trip_impl()) + })?; + + match handle.join() { + Ok(result) => result, + Err(_) => Err(anyhow::anyhow!( + "streamable_http_with_oauth_round_trip thread panicked" + )), + } +} + +#[allow(clippy::expect_used)] +async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + // Phase 1: script the model responses so Codex will call the OAuth-backed + // MCP echo tool and then finish the turn after receiving the result. + let server = responses::start_mock_server().await; + + let call_id = "call-789"; + let server_name = "rmcp_http_oauth"; + let tool_name = format!("mcp__{server_name}__echo"); + let namespace = format!("mcp__{server_name}__"); + + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call_with_namespace( + call_id, + &namespace, + "echo", + "{\"message\":\"ping\"}", + ), + responses::ev_completed("resp-1"), + ]), + ) + .await; + mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_assistant_message( + "msg-1", + "rmcp streamable http oauth echo tool completed successfully.", + ), + responses::ev_completed("resp-2"), + ]), + ) + .await; + + // Phase 2: start the Streamable HTTP MCP test server with bearer-token + // enforcement enabled so the client must use stored OAuth credentials. + let expected_env_value = "propagated-env-http-oauth"; + 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? + else { + return Ok(()); + }; + let server_url = http_server.url().to_string(); + + // Phase 3: seed an isolated CODEX_HOME with fallback OAuth tokens for this + // server so the test does not share credentials with other suite cases. + let temp_home = Arc::new(tempdir()?); + let _codex_home_guard = EnvVarGuard::set("CODEX_HOME", temp_home.path().as_os_str()); + write_fallback_oauth_tokens( + temp_home.path(), + server_name, + &server_url, + client_id, + expected_token, + refresh_token, + )?; + + // Phase 4: configure Codex with the OAuth-backed Streamable HTTP MCP + // server and build the fixture in the active local or remote-aware mode. + let fixture = test_codex() + .with_home(temp_home.clone()) + .with_config(move |config| { + // Keep OAuth credentials isolated to this test home because Bazel + // runs the full core suite in one process. + config.mcp_oauth_credentials_store_mode = serde_json::from_value(json!("file")) + .expect("`file` should deserialize as OAuthCredentialsStoreMode"); + insert_mcp_server( + config, + server_name, + McpServerTransportConfig::StreamableHttp { + url: server_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + TestMcpServerOptions { + experimental_environment: remote_aware_experimental_environment(), + ..Default::default() + }, + ); + }) + .build_remote_aware(&server) + .await?; + let session_model = fixture.session_configured.model.clone(); + + // Phase 5: wait for MCP discovery to publish the expected tool before the + // turn is submitted, which keeps failures tied to server startup/discovery. + wait_for_mcp_tool(&fixture, &tool_name).await?; + + // Phase 6: submit the user turn that should invoke the OAuth-backed tool. + fixture + .codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "call the rmcp streamable http oauth echo tool".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: fixture.cwd.path().to_path_buf(), + approval_policy: AskForApproval::Never, + approvals_reviewer: None, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: session_model, + effort: None, + summary: None, + service_tier: None, + collaboration_mode: None, + personality: None, + environments: None, + }) + .await?; + + // Phase 7: assert Codex begins the expected tool invocation. + let begin_event = wait_for_event(&fixture.codex, |ev| { + matches!(ev, EventMsg::McpToolCallBegin(_)) + }) + .await; + + let EventMsg::McpToolCallBegin(begin) = begin_event else { + unreachable!("event guard guarantees McpToolCallBegin"); + }; + assert_eq!(begin.invocation.server, server_name); + assert_eq!(begin.invocation.tool, "echo"); + + // Phase 8: assert the tool result proves the authenticated request reached + // the server and preserved the expected environment value. + let end_event = wait_for_event(&fixture.codex, |ev| { + matches!(ev, EventMsg::McpToolCallEnd(_)) + }) + .await; + let EventMsg::McpToolCallEnd(end) = end_event else { + unreachable!("event guard guarantees McpToolCallEnd"); + }; + + let result = end + .result + .as_ref() + .expect("rmcp echo tool should return success"); + assert_eq!(result.is_error, Some(false)); + assert!( + result.content.is_empty(), + "content should default to an empty array" + ); + + let structured = result + .structured_content + .as_ref() + .expect("structured content"); + let Value::Object(map) = structured else { + panic!("structured content should be an object: {structured:?}"); + }; + let echo_value = map + .get("echo") + .and_then(Value::as_str) + .expect("echo payload present"); + assert_eq!(echo_value, "ECHOING: ping"); + let env_value = map + .get("env") + .and_then(Value::as_str) + .expect("env snapshot inserted"); + assert_eq!(env_value, expected_env_value); + + // Phase 9: verify the scripted model calls were consumed and clean up the + // placement-aware MCP server. + wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + server.verify().await; + + http_server.shutdown().await; + + Ok(()) +} + +/// Starts the Streamable HTTP MCP test server in the active test placement. +async fn start_streamable_http_test_server( + expected_env_value: &str, + expected_token: Option<&str>, +) -> anyhow::Result> { + let rmcp_http_server_bin = match cargo_bin("test_streamable_http_server") { + Ok(path) => path, + Err(err) => { + eprintln!("test_streamable_http_server binary not available, skipping test: {err}"); + return Ok(None); + } + }; + + if let Some(container_name) = remote_env_container_name()? { + return Ok(Some( + start_remote_streamable_http_test_server( + &container_name, + &rmcp_http_server_bin, + expected_env_value, + expected_token, + ) + .await?, + )); + } + + let listener = TcpListener::bind("127.0.0.1:0")?; + let port = listener.local_addr()?.port(); + drop(listener); + let bind_addr = format!("127.0.0.1:{port}"); + let server_url = format!("http://{bind_addr}/mcp"); + + let mut command = Command::new(&rmcp_http_server_bin); + command + .kill_on_drop(true) + .env("MCP_STREAMABLE_HTTP_BIND_ADDR", &bind_addr) + .env("MCP_TEST_VALUE", expected_env_value); + if let Some(expected_token) = expected_token { + command.env("MCP_EXPECT_BEARER", expected_token); + } + let mut child = command.spawn()?; + + wait_for_local_streamable_http_server(&mut child, &server_url, Duration::from_secs(5)).await?; + Ok(Some(StreamableHttpTestServer { + server_url, + process: StreamableHttpTestServerProcess::Local(child), + })) +} + +/// Starts the Streamable HTTP MCP test server inside the remote test container. +async fn start_remote_streamable_http_test_server( + container_name: &str, + rmcp_http_server_bin: &Path, + expected_env_value: &str, + expected_token: Option<&str>, +) -> anyhow::Result { + let remote_path = copy_binary_to_remote_env( + container_name, + rmcp_http_server_bin, + "test_streamable_http_server", + )?; + let bound_addr_file = format!("{remote_path}.addr"); + let log_file = format!("{remote_path}.log"); + let mut env_assignments = vec![ + format!( + "MCP_STREAMABLE_HTTP_BIND_ADDR={}", + sh_single_quote("0.0.0.0:0") + ), + format!( + "MCP_STREAMABLE_HTTP_BOUND_ADDR_FILE={}", + sh_single_quote(&bound_addr_file) + ), + format!("MCP_TEST_VALUE={}", sh_single_quote(expected_env_value)), + ]; + if let Some(expected_token) = expected_token { + env_assignments.push(format!( + "MCP_EXPECT_BEARER={}", + sh_single_quote(expected_token) + )); + } + + let script = format!( + "{} nohup {} > {} 2>&1 < /dev/null & echo $!", + env_assignments.join(" "), + sh_single_quote(&remote_path), + sh_single_quote(&log_file) + ); + let start_output = StdCommand::new("docker") + .args(["exec", container_name, "sh", "-lc", &script]) + .output() + .context("start remote streamable HTTP MCP test server")?; + ensure!( + start_output.status.success(), + "docker start streamable HTTP MCP test server failed: stdout={} stderr={}", + String::from_utf8_lossy(&start_output.stdout).trim(), + String::from_utf8_lossy(&start_output.stderr).trim() + ); + let pid = String::from_utf8(start_output.stdout) + .context("remote streamable HTTP server pid must be utf-8")? + .trim() + .to_string(); + ensure!( + !pid.is_empty(), + "remote streamable HTTP server pid is empty" + ); + + let remote_bind_addr = + wait_for_remote_bound_addr(container_name, &bound_addr_file, Duration::from_secs(5)) + .await?; + let container_ip = remote_container_ip(container_name)?; + let server_url = format!("http://{}:{}/mcp", container_ip, remote_bind_addr.port()); + // The orchestrator can see the Docker container IP, but the behavior under + // test is whether the remote-side MCP client can reach it. Probe through + // remote HTTP before handing the URL to the Codex fixture. + wait_for_remote_streamable_http_server(&server_url, Duration::from_secs(5)).await?; + if expected_token.is_some() { + wait_for_streamable_http_metadata(&server_url, Duration::from_secs(5)).await?; + } + + Ok(StreamableHttpTestServer { + server_url, + process: StreamableHttpTestServerProcess::Remote(RemoteStreamableHttpServer { + container_name: container_name.to_string(), + pid, + paths_to_remove: vec![remote_path, bound_addr_file, log_file], + }), + }) +} + +/// Single-quotes a value for the small shell snippets sent through Docker. +fn sh_single_quote(value: &str) -> String { + format!("'{}'", value.replace('\'', "'\\''")) +} + +/// Waits until the remote test server writes the socket address it bound to. +async fn wait_for_remote_bound_addr( + container_name: &str, + bound_addr_file: &str, + timeout: Duration, +) -> anyhow::Result { + let deadline = Instant::now() + timeout; + loop { + let output = StdCommand::new("docker") + .args(["exec", container_name, "cat", bound_addr_file]) + .output() + .context("read remote streamable HTTP server bound address")?; + if output.status.success() { + let bound_addr = String::from_utf8(output.stdout) + .context("remote streamable HTTP bound address must be utf-8")?; + return bound_addr + .trim() + .parse() + .context("parse remote streamable HTTP bound address"); + } + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for remote streamable HTTP bound address: {}", + String::from_utf8_lossy(&output.stderr).trim() + )); + } + sleep(Duration::from_millis(50)).await; + } +} + +/// Reads the container IP that the host-side test process can use. +fn remote_container_ip(container_name: &str) -> anyhow::Result { + let output = StdCommand::new("docker") + .args([ + "inspect", + "-f", + "{{range .NetworkSettings.Networks}}{{println .IPAddress}}{{end}}", + container_name, + ]) + .output() + .context("inspect remote MCP test container IP")?; + ensure!( + output.status.success(), + "docker inspect remote MCP test container IP failed: stdout={} stderr={}", + String::from_utf8_lossy(&output.stdout).trim(), + String::from_utf8_lossy(&output.stderr).trim() + ); + let inspect_output = + String::from_utf8(output.stdout).context("remote MCP test container IP must be utf-8")?; + let ip = inspect_output + .lines() + .map(str::trim) + .find(|line| !line.is_empty()) + .unwrap_or_default() + .to_string(); + if ip.is_empty() { + Ok("127.0.0.1".to_string()) + } else { + Ok(ip) + } +} + +/// Waits for the local Streamable HTTP test server to publish OAuth metadata. +async fn wait_for_local_streamable_http_server( + server_child: &mut Child, + server_url: &str, + timeout: Duration, +) -> anyhow::Result<()> { + let deadline = Instant::now() + timeout; + let metadata_url = streamable_http_metadata_url(server_url); + let client = Client::builder().no_proxy().build()?; + loop { + if let Some(status) = server_child.try_wait()? { + return Err(anyhow::anyhow!( + "streamable HTTP server exited early with status {status}" + )); + } + + let remaining = deadline.saturating_duration_since(Instant::now()); + + if remaining.is_zero() { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: deadline reached" + )); + } + + match tokio::time::timeout(remaining, client.get(&metadata_url).send()).await { + Ok(Ok(response)) if response.status() == StatusCode::OK => return Ok(()), + Ok(Ok(response)) => { + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: HTTP {}", + response.status() + )); + } + } + Ok(Err(error)) => { + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: {error}" + )); + } + } + Err(_) => { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: request timed out" + )); + } + } + + sleep(Duration::from_millis(50)).await; + } +} + +/// Waits for the remote Streamable HTTP test server via remote HTTP. +async fn wait_for_remote_streamable_http_server( + server_url: &str, + timeout: Duration, +) -> anyhow::Result<()> { + let websocket_url = std::env::var(REMOTE_EXEC_SERVER_URL_ENV_VAR).with_context(|| { + format!("{REMOTE_EXEC_SERVER_URL_ENV_VAR} must be set for remote streamable HTTP MCP tests") + })?; + let environment = Environment::create_for_tests(Some(websocket_url))?; + let http_client = environment.get_http_client(); + let metadata_url = streamable_http_metadata_url(server_url); + let deadline = Instant::now() + timeout; + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(anyhow::anyhow!( + "timed out waiting for remote streamable HTTP server metadata at {metadata_url}: deadline reached" + )); + } + + let request = HttpRequestParams { + method: "GET".to_string(), + url: metadata_url.clone(), + headers: Vec::new(), + body: None, + timeout_ms: Some(remaining.as_millis().clamp(1, 1_000) as u64), + request_id: "buffered-request".to_string(), + stream_response: false, + }; + match http_client.http_request(request).await { + Ok(response) if response.status == StatusCode::OK.as_u16() => return Ok(()), + Ok(response) => { + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for remote streamable HTTP server metadata at {metadata_url}: HTTP {}", + response.status + )); + } + } + Err(error) => { + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for remote streamable HTTP server metadata at {metadata_url}: {error}" + )); + } + } + } + + sleep(Duration::from_millis(50)).await; + } +} + +/// Waits for OAuth metadata from the host-side test process. +async fn wait_for_streamable_http_metadata( + server_url: &str, + timeout: Duration, +) -> anyhow::Result<()> { + let deadline = Instant::now() + timeout; + let metadata_url = streamable_http_metadata_url(server_url); + let client = Client::builder().no_proxy().build()?; + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: deadline reached" + )); + } + + match tokio::time::timeout(remaining, client.get(&metadata_url).send()).await { + Ok(Ok(response)) if response.status() == StatusCode::OK => return Ok(()), + Ok(Ok(response)) => { + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: HTTP {}", + response.status() + )); + } + } + Ok(Err(error)) => { + if Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: {error}" + )); + } + } + Err(_) => { + return Err(anyhow::anyhow!( + "timed out waiting for streamable HTTP server metadata at {metadata_url}: request timed out" + )); + } + } + + sleep(Duration::from_millis(50)).await; + } +} + +/// Builds the OAuth metadata URL for the test Streamable HTTP MCP endpoint. +fn streamable_http_metadata_url(server_url: &str) -> String { + let base_url = server_url.strip_suffix("/mcp").unwrap_or(server_url); + format!("{base_url}{STREAMABLE_HTTP_METADATA_PATH}") +} 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 b948cab461..d1c22f430c 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 @@ -1,5 +1,6 @@ use std::borrow::Cow; use std::collections::HashMap; +use std::fs; use std::io::ErrorKind; use std::net::SocketAddr; use std::sync::Arc; @@ -9,11 +10,13 @@ use axum::Router; use axum::body::Body; use axum::extract::Json; use axum::extract::State; +use axum::http::HeaderMap; use axum::http::Method; use axum::http::Request; use axum::http::StatusCode; use axum::http::header::AUTHORIZATION; use axum::http::header::CONTENT_TYPE; +use axum::http::header::HOST; use axum::middleware; use axum::middleware::Next; use axum::response::Response; @@ -60,90 +63,6 @@ const MEMO_CONTENT: &str = "This is a sample MCP resource served by the rmcp tes const MCP_SESSION_ID_HEADER: &str = "mcp-session-id"; const SESSION_POST_FAILURE_CONTROL_PATH: &str = "/test/control/session-post-failure"; -impl TestToolServer { - fn new() -> Self { - let tools = vec![Self::echo_tool()]; - let resources = vec![Self::memo_resource()]; - let resource_templates = vec![Self::memo_template()]; - Self { - tools: Arc::new(tools), - resources: Arc::new(resources), - resource_templates: Arc::new(resource_templates), - } - } - - fn echo_tool() -> Tool { - #[expect(clippy::expect_used)] - let schema: JsonObject = serde_json::from_value(json!({ - "type": "object", - "properties": { - "message": { "type": "string" }, - "env_var": { "type": "string" } - }, - "required": ["message"], - "additionalProperties": false - })) - .expect("echo tool schema should deserialize"); - - let mut tool = Tool::new( - Cow::Borrowed("echo"), - Cow::Borrowed("Echo back the provided message and include environment data."), - Arc::new(schema), - ); - #[expect(clippy::expect_used)] - let output_schema: JsonObject = serde_json::from_value(json!({ - "type": "object", - "properties": { - "echo": { "type": "string" }, - "env": { - "anyOf": [ - { "type": "string" }, - { "type": "null" } - ] - } - }, - "required": ["echo", "env"], - "additionalProperties": false - })) - .expect("echo tool output schema should deserialize"); - tool.output_schema = Some(Arc::new(output_schema)); - tool.annotations = Some(ToolAnnotations::new().read_only(true)); - tool - } - - fn memo_resource() -> Resource { - let raw = RawResource { - uri: MEMO_URI.to_string(), - name: "example-note".to_string(), - title: Some("Example Note".to_string()), - description: Some("A sample MCP resource exposed for integration tests.".to_string()), - mime_type: Some("text/plain".to_string()), - size: None, - icons: None, - meta: None, - }; - Resource::new(raw, None) - } - - fn memo_template() -> ResourceTemplate { - let raw = RawResourceTemplate { - uri_template: "memo://codex/{slug}".to_string(), - name: "codex-memo".to_string(), - title: Some("Codex Memo".to_string()), - description: Some( - "Template for memo://codex/{slug} resources used in tests.".to_string(), - ), - mime_type: Some("text/plain".to_string()), - icons: None, - }; - ResourceTemplate::new(raw, None) - } - - fn memo_text() -> &'static str { - MEMO_CONTENT - } -} - #[derive(Clone, Default)] struct SessionFailureState { armed_failure: Arc>>, @@ -168,6 +87,91 @@ struct EchoArgs { env_var: Option, } +#[tokio::main] +async fn main() -> Result<(), Box> { + let bind_addr = parse_bind_addr()?; + let session_failure_state = SessionFailureState::default(); + const MAX_BIND_RETRIES: u32 = 20; + const BIND_RETRY_DELAY: Duration = Duration::from_millis(50); + + let mut bind_retries = 0; + let listener = loop { + match tokio::net::TcpListener::bind(&bind_addr).await { + Ok(listener) => break listener, + Err(err) if err.kind() == ErrorKind::PermissionDenied => { + eprintln!( + "failed to bind to {bind_addr}: {err}. make sure the process has network access" + ); + return Ok(()); + } + Err(err) if err.kind() == ErrorKind::AddrInUse && bind_retries < MAX_BIND_RETRIES => { + bind_retries += 1; + sleep(BIND_RETRY_DELAY).await; + } + Err(err) => return Err(err.into()), + } + }; + let actual_bind_addr = listener.local_addr()?; + if let Ok(bound_addr_file) = std::env::var("MCP_STREAMABLE_HTTP_BOUND_ADDR_FILE") { + fs::write(bound_addr_file, actual_bind_addr.to_string())?; + } + eprintln!("starting rmcp streamable http test server on http://{actual_bind_addr}/mcp"); + + let router = Router::new() + .route( + SESSION_POST_FAILURE_CONTROL_PATH, + post(arm_session_post_failure), + ) + .route( + "/.well-known/oauth-authorization-server/mcp", + get({ + move |headers: HeaderMap| async move { + let metadata_base = headers + .get(HOST) + .and_then(|value| value.to_str().ok()) + .map(|host| format!("http://{host}")) + .unwrap_or_else(|| format!("http://{actual_bind_addr}")); + #[expect(clippy::expect_used)] + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + serde_json::to_vec(&json!({ + "authorization_endpoint": format!("{metadata_base}/oauth/authorize"), + "token_endpoint": format!("{metadata_base}/oauth/token"), + "scopes_supported": [""], + })).expect("failed to serialize metadata"), + )) + .expect("valid metadata response") + } + }), + ) + .nest_service( + "/mcp", + StreamableHttpService::new( + || Ok(TestToolServer::new()), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ), + ) + .layer(middleware::from_fn_with_state( + session_failure_state.clone(), + fail_session_post_when_armed, + )) + .with_state(session_failure_state); + + let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") { + let expected = Arc::new(format!("Bearer {token}")); + router.layer(middleware::from_fn_with_state(expected, require_bearer)) + } else { + router + }; + + axum::serve(listener, router).await?; + task::yield_now().await; + Ok(()) +} + impl ServerHandler for TestToolServer { fn get_info(&self) -> ServerInfo { ServerInfo { @@ -285,6 +289,90 @@ impl ServerHandler for TestToolServer { } } +impl TestToolServer { + fn new() -> Self { + let tools = vec![Self::echo_tool()]; + let resources = vec![Self::memo_resource()]; + let resource_templates = vec![Self::memo_template()]; + Self { + tools: Arc::new(tools), + resources: Arc::new(resources), + resource_templates: Arc::new(resource_templates), + } + } + + fn echo_tool() -> Tool { + #[expect(clippy::expect_used)] + let schema: JsonObject = serde_json::from_value(json!({ + "type": "object", + "properties": { + "message": { "type": "string" }, + "env_var": { "type": "string" } + }, + "required": ["message"], + "additionalProperties": false + })) + .expect("echo tool schema should deserialize"); + + let mut tool = Tool::new( + Cow::Borrowed("echo"), + Cow::Borrowed("Echo back the provided message and include environment data."), + Arc::new(schema), + ); + #[expect(clippy::expect_used)] + let output_schema: JsonObject = serde_json::from_value(json!({ + "type": "object", + "properties": { + "echo": { "type": "string" }, + "env": { + "anyOf": [ + { "type": "string" }, + { "type": "null" } + ] + } + }, + "required": ["echo", "env"], + "additionalProperties": false + })) + .expect("echo tool output schema should deserialize"); + tool.output_schema = Some(Arc::new(output_schema)); + tool.annotations = Some(ToolAnnotations::new().read_only(true)); + tool + } + + fn memo_resource() -> Resource { + let raw = RawResource { + uri: MEMO_URI.to_string(), + name: "example-note".to_string(), + title: Some("Example Note".to_string()), + description: Some("A sample MCP resource exposed for integration tests.".to_string()), + mime_type: Some("text/plain".to_string()), + size: None, + icons: None, + meta: None, + }; + Resource::new(raw, None) + } + + fn memo_template() -> ResourceTemplate { + let raw = RawResourceTemplate { + uri_template: "memo://codex/{slug}".to_string(), + name: "codex-memo".to_string(), + title: Some("Codex Memo".to_string()), + description: Some( + "Template for memo://codex/{slug} resources used in tests.".to_string(), + ), + mime_type: Some("text/plain".to_string()), + icons: None, + }; + ResourceTemplate::new(raw, None) + } + + fn memo_text() -> &'static str { + MEMO_CONTENT + } +} + fn parse_bind_addr() -> Result> { let default_addr = "127.0.0.1:3920"; let bind_addr = std::env::var("MCP_STREAMABLE_HTTP_BIND_ADDR") @@ -293,83 +381,6 @@ fn parse_bind_addr() -> Result> { Ok(bind_addr.parse()?) } -#[tokio::main] -async fn main() -> Result<(), Box> { - let bind_addr = parse_bind_addr()?; - let session_failure_state = SessionFailureState::default(); - const MAX_BIND_RETRIES: u32 = 20; - const BIND_RETRY_DELAY: Duration = Duration::from_millis(50); - - let mut bind_retries = 0; - let listener = loop { - match tokio::net::TcpListener::bind(&bind_addr).await { - Ok(listener) => break listener, - Err(err) if err.kind() == ErrorKind::PermissionDenied => { - eprintln!( - "failed to bind to {bind_addr}: {err}. make sure the process has network access" - ); - return Ok(()); - } - Err(err) if err.kind() == ErrorKind::AddrInUse && bind_retries < MAX_BIND_RETRIES => { - bind_retries += 1; - sleep(BIND_RETRY_DELAY).await; - } - Err(err) => return Err(err.into()), - } - }; - eprintln!("starting rmcp streamable http test server on http://{bind_addr}/mcp"); - - let router = Router::new() - .route( - SESSION_POST_FAILURE_CONTROL_PATH, - post(arm_session_post_failure), - ) - .route( - "/.well-known/oauth-authorization-server/mcp", - get({ - move || async move { - let metadata_base = format!("http://{bind_addr}"); - #[expect(clippy::expect_used)] - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "application/json") - .body(Body::from( - serde_json::to_vec(&json!({ - "authorization_endpoint": format!("{metadata_base}/oauth/authorize"), - "token_endpoint": format!("{metadata_base}/oauth/token"), - "scopes_supported": [""], - })).expect("failed to serialize metadata"), - )) - .expect("valid metadata response") - } - }), - ) - .nest_service( - "/mcp", - StreamableHttpService::new( - || Ok(TestToolServer::new()), - Arc::new(LocalSessionManager::default()), - StreamableHttpServerConfig::default(), - ), - ) - .layer(middleware::from_fn_with_state( - session_failure_state.clone(), - fail_session_post_when_armed, - )) - .with_state(session_failure_state); - - let router = if let Ok(token) = std::env::var("MCP_EXPECT_BEARER") { - let expected = Arc::new(format!("Bearer {token}")); - router.layer(middleware::from_fn_with_state(expected, require_bearer)) - } else { - router - }; - - axum::serve(listener, router).await?; - task::yield_now().await; - Ok(()) -} - async fn require_bearer( State(expected): State>, request: Request,