diff --git a/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs new file mode 100644 index 0000000000..3f6afd99f3 --- /dev/null +++ b/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs @@ -0,0 +1,250 @@ +use std::io; +use std::io::BufRead; +use std::io::Write; + +use anyhow::Result; +use serde_json::Value; +use serde_json::json; + +const MODERN_VERSION: &str = "2026-07-28"; +const LEGACY_VERSION: &str = "2025-06-18"; +const MAX_MESSAGE_BYTES: usize = 8 * 1024 * 1024; + +fn write_message(stdout: &mut io::StdoutLock<'_>, message: Value) -> Result<()> { + serde_json::to_writer(&mut *stdout, &message)?; + stdout.write_all(b"\n")?; + stdout.flush()?; + Ok(()) +} + +fn assert_modern_metadata(message: &Value) { + let metadata = &message["params"]["_meta"]; + assert_eq!( + metadata["io.modelcontextprotocol/protocolVersion"], + MODERN_VERSION + ); + assert!(metadata["io.modelcontextprotocol/clientInfo"].is_object()); + assert!(metadata["io.modelcontextprotocol/clientCapabilities"].is_object()); +} + +fn legacy_schema_defaults() -> Value { + json!({ + "name": "John Doe", + "age": 30, + "score": 95.5, + "status": "active", + "verified": true, + }) +} + +fn main() -> Result<()> { + let mode = std::env::args().nth(1).unwrap_or_default(); + assert!(matches!( + mode.as_str(), + "modern" | "legacy" | "legacy-fallback" | "oversized-stdout" | "oversized-stdout-legacy" + )); + assert!(std::env::var_os("CODEX_MCP_PROTOCOL_VERSION").is_none()); + let is_legacy = mode.starts_with("legacy") || mode.ends_with("-legacy"); + let oversized = mode.starts_with("oversized-stdout"); + let mut stdout = io::stdout().lock(); + let mut initialized = false; + let mut pending_legacy_call = None; + + for line in io::stdin().lock().lines() { + let message: Value = serde_json::from_str(&line?)?; + let Some(method) = message.get("method").and_then(Value::as_str) else { + let pending_call = pending_legacy_call + .take() + .ok_or_else(|| anyhow::anyhow!("unexpected server-request response"))?; + assert_eq!(message["id"], "legacy-approval"); + assert_eq!(message["result"]["action"], "accept"); + assert_eq!(message["result"]["content"], legacy_schema_defaults()); + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": pending_call, + "result": {"content": [{"type": "text", "text": "legacy approved"}]}, + }), + )?; + continue; + }; + + match method { + "server/discover" if is_legacy => { + assert_eq!(mode, "legacy-fallback"); + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": message["id"], + "error": {"code": -32601, "message": "method not found"}, + }), + )?; + } + "server/discover" => { + assert!(!initialized); + assert_modern_metadata(&message); + initialized = true; + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "resultType": "complete", + "supportedVersions": [MODERN_VERSION], + "capabilities": {"tools": {}, "resources": {}}, + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "strict-stdio-test", + "version": "1.0.0", + }, + }, + "ttlMs": 0, + "cacheScope": "private", + }, + }), + )?; + } + "initialize" => { + assert!(is_legacy); + assert_eq!(message["params"]["protocolVersion"], LEGACY_VERSION); + initialized = true; + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "protocolVersion": LEGACY_VERSION, + "capabilities": {"tools": {}}, + "serverInfo": {"name": "legacy-stdio-test", "version": "1.0.0"}, + }, + }), + )?; + } + "notifications/initialized" => assert!(is_legacy), + "tools/list" => { + assert!(initialized); + if !is_legacy { + assert_modern_metadata(&message); + } + let description = if oversized { + "x".repeat(MAX_MESSAGE_BYTES + 1) + } else { + "echo a value".to_owned() + }; + let mut result = json!({ + "tools": [{ + "name": "echo", + "description": description, + "inputSchema": {"type": "object"}, + }], + }); + if !is_legacy { + result["resultType"] = json!("complete"); + result["ttlMs"] = json!(0); + result["cacheScope"] = json!("private"); + } + write_message( + &mut stdout, + json!({"jsonrpc": "2.0", "id": message["id"], "result": result}), + )?; + } + "tools/call" if is_legacy => { + assert!(initialized); + pending_legacy_call = Some(message["id"].clone()); + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": "legacy-approval", + "method": "elicitation/create", + "params": { + "message": "Accept the schema defaults", + "requestedSchema": { + "type": "object", + "properties": { + "name": {"type": "string", "default": "John Doe"}, + "age": {"type": "integer", "default": 30}, + "score": {"type": "number", "default": 95.5}, + "status": { + "type": "string", + "enum": ["active", "inactive"], + "default": "active", + }, + "verified": {"type": "boolean", "default": true}, + }, + "required": [], + }, + }, + }), + )?; + } + "tools/call" => { + assert!(initialized); + assert_modern_metadata(&message); + if message.pointer("/params/inputResponses").is_none() { + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "resultType": "input_required", + "inputRequests": { + "approval": { + "method": "elicitation/create", + "params": { + "mode": "form", + "message": "Approve stdio call?", + "requestedSchema": { + "type": "object", + "properties": { + "approved": {"type": "boolean"}, + }, + }, + }, + }, + }, + "requestState": "stdio-state", + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "strict-stdio-test", + "version": "1.0.0", + }, + }, + }, + }), + )?; + continue; + } + + assert_eq!(message["params"]["requestState"], "stdio-state"); + assert_eq!( + message["params"]["inputResponses"]["approval"]["action"], + "accept" + ); + assert_eq!( + message["params"]["inputResponses"]["approval"]["content"]["approved"], + true + ); + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", + "id": message["id"], + "result": { + "resultType": "complete", + "content": [{"type": "text", "text": "modern approved"}], + }, + }), + )?; + } + other => anyhow::bail!("unexpected MCP method {other}"), + } + } + + Ok(()) +} diff --git a/codex-rs/rmcp-client/src/elicitation_client_service.rs b/codex-rs/rmcp-client/src/elicitation_client_service.rs index 5d89c08ca9..84cef31d04 100644 --- a/codex-rs/rmcp-client/src/elicitation_client_service.rs +++ b/codex-rs/rmcp-client/src/elicitation_client_service.rs @@ -5,7 +5,10 @@ use rmcp::model::ClientInfo; use rmcp::model::ClientResult; use rmcp::model::CustomRequest; use rmcp::model::CustomResult; +use rmcp::model::ElicitResult; use rmcp::model::ElicitationAction; +use rmcp::model::MetaObject; +use rmcp::model::ProtocolVersion; use rmcp::model::RequestMetaObject; use rmcp::model::RequestParamsMeta; use rmcp::model::ServerNotification; @@ -94,22 +97,42 @@ impl Service for ElicitationClientService { ) -> Result { match request { ServerRequest::ElicitRequest(request) => { + let modern_session = context + .peer + .peer_info() + .is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); let response = self .create_elicitation(Elicitation::Mcp(request.params), context) .await?; - // RMCP's typed ElicitResult does not model result-level `_meta`. - let result = elicitation_response_result(response)?; - Ok(ClientResult::CustomResult(result)) + if modern_session { + Ok(ClientResult::ElicitResult(typed_elicitation_result( + response, + )?)) + } else { + Ok(ClientResult::CustomResult(elicitation_response_result( + response, + )?)) + } } ServerRequest::CustomRequest(request) if request.method == MCP_ELICITATION_CREATE_METHOD => { + let modern_session = context + .peer + .peer_info() + .is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); let response = self .create_elicitation(custom_mcp_elicitation(request)?, context) .await?; - Ok(ClientResult::CustomResult(elicitation_response_result( - response, - )?)) + if modern_session { + Ok(ClientResult::ElicitResult(typed_elicitation_result( + response, + )?)) + } else { + Ok(ClientResult::CustomResult(elicitation_response_result( + response, + )?)) + } } ServerRequest::CustomRequest(request) if request.method == OPENAI_FORM_METHOD && self.supports_openai_form => @@ -227,6 +250,29 @@ fn elicitation_response_result( .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None)) } +fn typed_elicitation_result( + response: ElicitationResponse, +) -> Result { + let ElicitationResponse { + action, + content, + meta, + } = response; + let mut result = ElicitResult::new(action); + result.content = content; + result.meta = match meta { + None => None, + Some(Value::Object(meta)) => Some(MetaObject::from(meta)), + Some(meta) => { + return Err(rmcp::ErrorData::invalid_params( + format!("MCP elicitation response _meta must be an object, got {meta}"), + None, + )); + } + }; + Ok(result) +} + #[cfg(test)] mod tests { use pretty_assertions::assert_eq; @@ -286,7 +332,7 @@ mod tests { let ServerRequest::ElicitRequest(request) = request else { panic!("legacy elicitation/create must dispatch to the typed handler"); }; - let rmcp::model::ElicitRequestParams::FormElicitationParams { + let ElicitRequestParams::FormElicitationParams { requested_schema, .. } = request.params else { @@ -418,6 +464,37 @@ mod tests { ); } + #[test] + fn typed_elicitation_result_preserves_response_meta() { + let result = typed_elicitation_result(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({ "confirmed": true })), + meta: Some(json!({ "persist": "always" })), + }) + .expect("modern elicitation response should serialize"); + + assert_eq!( + serde_json::to_value(result).expect("typed elicitation result should serialize"), + json!({ + "action": "accept", + "content": { "confirmed": true }, + "_meta": { "persist": "always" }, + }) + ); + } + + #[test] + fn typed_elicitation_result_rejects_non_object_meta() { + let error = typed_elicitation_result(ElicitationResponse { + action: ElicitationAction::Accept, + content: None, + meta: Some(json!(["invalid"])), + }) + .expect_err("modern elicitation metadata must be an object"); + + assert_eq!(error.code, rmcp::model::ErrorCode::INVALID_PARAMS); + } + fn form_request(meta: Option) -> ElicitRequestParams { ElicitRequestParams::FormElicitationParams { meta, diff --git a/codex-rs/rmcp-client/src/executor_process_transport.rs b/codex-rs/rmcp-client/src/executor_process_transport.rs index 4f3271f6e1..f36c14f70e 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport.rs @@ -356,11 +356,24 @@ impl ExecutorProcessTransport { .await .map_err(io::Error::other)?; for chunk in response.chunks { + let expected_seq = self.last_seq.saturating_add(1); + if chunk.seq > expected_seq { + return Err(self.close_for_lost_output(expected_seq, chunk.seq)); + } self.push_process_output_if_new(chunk); if self.closed { return Ok(()); } } + // Process reads include output chunks but not the sequenced `Exited` + // and `Closed` events. Account for those terminal events without + // allowing an evicted output chunk to be silently spliced into MCP. + let terminal_event_count = u64::from(response.exited) + u64::from(response.closed); + let next_output_seq = response.next_seq.saturating_sub(terminal_event_count); + let expected_seq = self.last_seq.saturating_add(1); + if next_output_seq > expected_seq { + return Err(self.close_for_lost_output(expected_seq, next_output_seq)); + } self.last_seq = self.last_seq.max(response.next_seq.saturating_sub(1)); if let Some(message) = response.failure { warn!( @@ -374,6 +387,18 @@ impl ExecutorProcessTransport { Ok(()) } + fn close_for_lost_output(&mut self, expected_seq: u64, received_seq: u64) -> io::Error { + self.stdout.clear(); + self.stderr.clear(); + self.closed = true; + io::Error::new( + io::ErrorKind::InvalidData, + format!( + "remote MCP server output stream lost process events: expected sequence {expected_seq}, received {received_seq}" + ), + ) + } + fn push_process_output_if_new(&mut self, chunk: ProcessOutputChunk) { if !self.should_accept_seq(chunk.seq) { return; diff --git a/codex-rs/rmcp-client/src/executor_process_transport_tests.rs b/codex-rs/rmcp-client/src/executor_process_transport_tests.rs index 5b5b7e5e82..82fdaddcc9 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport_tests.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport_tests.rs @@ -1,8 +1,10 @@ use bytes::BytesMut; +use codex_exec_server::ExecOutputStream; use codex_exec_server::ExecProcess; use codex_exec_server::ExecProcessEventReceiver; use codex_exec_server::ExecProcessFuture; use codex_exec_server::ProcessId; +use codex_exec_server::ProcessOutputChunk; use codex_exec_server::ProcessSignal; use codex_exec_server::ReadResponse; use codex_exec_server::WriteResponse; @@ -12,6 +14,7 @@ use rmcp::service::RoleClient; use rmcp::service::TxJsonRpcMessage; use rmcp::transport::Transport; use serde_json::json; +use std::io; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; @@ -22,6 +25,7 @@ use tokio::sync::watch; use super::ExecutorProcessTransport; use super::LineBuffer; use super::LineTooLong; +use super::MAX_MCP_STDERR_LINE_BYTES; use super::MAX_MCP_STDOUT_LINE_BYTES; struct BlockingFirstWriteProcess { @@ -83,6 +87,66 @@ impl ExecProcess for BlockingFirstWriteProcess { } } +struct RetainedReadProcess { + process_id: ProcessId, + response: ReadResponse, +} + +impl ExecProcess for RetainedReadProcess { + fn process_id(&self) -> &ProcessId { + &self.process_id + } + + fn subscribe_wake(&self) -> watch::Receiver { + watch::channel(0).1 + } + + fn subscribe_events(&self) -> ExecProcessEventReceiver { + ExecProcessEventReceiver::empty() + } + + fn read( + &self, + _after_seq: Option, + _max_bytes: Option, + _wait_ms: Option, + ) -> ExecProcessFuture<'_, ReadResponse> { + Box::pin(async { Ok(self.response.clone()) }) + } + + fn write(&self, _chunk: Vec) -> ExecProcessFuture<'_, WriteResponse> { + Box::pin(async { unreachable!("recovery tests should not write process input") }) + } + + fn signal(&self, _signal: ProcessSignal) -> ExecProcessFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } + + fn terminate(&self) -> ExecProcessFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } +} + +fn recovered_output(seq: u64, stream: ExecOutputStream, bytes: &[u8]) -> ProcessOutputChunk { + ProcessOutputChunk { + seq, + stream, + chunk: bytes.to_vec().into(), + } +} + +fn retained_read_response(chunks: Vec, next_seq: u64) -> ReadResponse { + ReadResponse { + chunks, + next_seq, + exited: false, + exit_code: None, + closed: false, + failure: None, + sandbox_denied: false, + } +} + #[tokio::test] async fn serializes_concurrent_stdin_writes() { let process = Arc::new(BlockingFirstWriteProcess { @@ -137,6 +201,177 @@ async fn serializes_concurrent_stdin_writes() { ); } +#[tokio::test] +async fn rejects_lagged_recovery_when_retained_output_skips_a_sequence() { + let process = Arc::new(RetainedReadProcess { + process_id: ProcessId::from("mcp-stdio-lost-output"), + response: retained_read_response( + vec![recovered_output( + /*seq*/ 6, + ExecOutputStream::Stdout, + br#"{"jsonrpc":"2.0","id":1,"result":{}}"#, + )], + /*next_seq*/ 7, + ), + }); + let mut transport = ExecutorProcessTransport::new(process, "mcp-stdio-lost-output".to_string()); + transport.last_seq = 4; + transport + .stdout + .extend_from_slice(br#"{"jsonrpc":"2.0","id":1,"result":""#) + .expect("partial stdout should fit"); + transport + .stderr + .extend_from_slice(b"partial diagnostic") + .expect("partial stderr should fit"); + + let error = transport + .recover_lagged_events() + .await + .expect_err("evicted output must close the transport"); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!( + error.to_string(), + "remote MCP server output stream lost process events: expected sequence 5, received 6" + ); + assert_eq!(transport.stdout, LineBuffer::default()); + assert_eq!(transport.stderr, LineBuffer::new(MAX_MCP_STDERR_LINE_BYTES)); + assert!(transport.closed); +} + +#[tokio::test] +async fn rejects_lagged_recovery_when_all_missing_output_was_evicted() { + let process = Arc::new(RetainedReadProcess { + process_id: ProcessId::from("mcp-stdio-evicted-output"), + response: retained_read_response(Vec::new(), /*next_seq*/ 7), + }); + let mut transport = + ExecutorProcessTransport::new(process, "mcp-stdio-evicted-output".to_string()); + transport.last_seq = 4; + + let error = transport + .recover_lagged_events() + .await + .expect_err("missing retained output must close the transport"); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert_eq!( + error.to_string(), + "remote MCP server output stream lost process events: expected sequence 5, received 7" + ); + assert!(transport.closed); +} + +#[tokio::test] +async fn recovers_contiguous_output_after_a_duplicate_replayed_event() { + let process = Arc::new(RetainedReadProcess { + process_id: ProcessId::from("mcp-stdio-duplicate-output"), + response: retained_read_response( + vec![ + recovered_output( + /*seq*/ 4, + ExecOutputStream::Stdout, + b"duplicate output\n", + ), + recovered_output( + /*seq*/ 5, + ExecOutputStream::Stdout, + b"recovered stdout\n", + ), + recovered_output( + /*seq*/ 6, + ExecOutputStream::Stderr, + b"recovered stderr\n", + ), + ], + /*next_seq*/ 7, + ), + }); + let mut transport = + ExecutorProcessTransport::new(process, "mcp-stdio-duplicate-output".to_string()); + transport.last_seq = 4; + + transport + .recover_lagged_events() + .await + .expect("duplicate replay must not prevent recovery of contiguous output"); + + assert_eq!( + transport.stdout.take_line(), + Some(BytesMut::from(&b"recovered stdout"[..])) + ); + assert_eq!(transport.stdout.take_line(), None); + assert_eq!(transport.stderr, LineBuffer::new(MAX_MCP_STDERR_LINE_BYTES)); + assert_eq!(transport.last_seq, 6); + assert!(!transport.closed); +} + +#[tokio::test] +async fn recovers_contiguous_output_and_accounts_for_terminal_events() { + let mut response = retained_read_response( + vec![ + recovered_output( + /*seq*/ 5, + ExecOutputStream::Stdout, + b"recovered stdout\n", + ), + recovered_output( + /*seq*/ 6, + ExecOutputStream::Stderr, + b"recovered stderr\n", + ), + ], + /*next_seq*/ 9, + ); + response.exited = true; + response.exit_code = Some(0); + response.closed = true; + let process = Arc::new(RetainedReadProcess { + process_id: ProcessId::from("mcp-stdio-recovered-output"), + response, + }); + let mut transport = + ExecutorProcessTransport::new(process, "mcp-stdio-recovered-output".to_string()); + transport.last_seq = 4; + + transport + .recover_lagged_events() + .await + .expect("contiguous retained output should remain recoverable"); + + assert_eq!( + transport.stdout.take_line(), + Some(BytesMut::from(&b"recovered stdout"[..])) + ); + assert_eq!(transport.stderr, LineBuffer::new(MAX_MCP_STDERR_LINE_BYTES)); + assert_eq!(transport.last_seq, 8); + assert!(transport.closed); +} + +#[tokio::test] +async fn closes_on_an_unsequenced_process_failure_without_inventing_missing_output() { + let mut response = retained_read_response(Vec::new(), /*next_seq*/ 5); + response.exited = true; + response.closed = true; + response.failure = Some("executor disconnected".to_string()); + let process = Arc::new(RetainedReadProcess { + process_id: ProcessId::from("mcp-stdio-failed-process"), + response, + }); + let mut transport = + ExecutorProcessTransport::new(process, "mcp-stdio-failed-process".to_string()); + transport.last_seq = 4; + + transport + .recover_lagged_events() + .await + .expect("an unsequenced executor failure is not an output sequence gap"); + + assert_eq!(transport.last_seq, 4); + assert!(transport.closed); +} + #[test] fn searches_only_new_bytes_after_partial_line() { let mut buffer = LineBuffer::default(); diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 87995b5bea..db0c579a9b 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -48,6 +48,7 @@ use sse_stream::Sse; use sse_stream::SseStream; use crate::incoming_jsonrpc::deserialize_incoming_jsonrpc_message; +use crate::incoming_jsonrpc::normalize_sse_jsonrpc_message; use crate::local_stdio_transport::MAX_MCP_STDIO_LINE_BYTES; mod www_authenticate; @@ -261,7 +262,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } match content_type.as_deref() { Some(content_type) if content_type.starts_with(EVENT_STREAM_MIME_TYPE) => { - let mut event_stream = sse_stream_from_body(body_stream); + let mut event_stream = sse_stream_from_body(body_stream, maximum_response_bytes); if mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE) { while let Some(event) = event_stream.next().await { let event = event.map_err(StreamableHttpError::Sse)?; @@ -292,7 +293,6 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { "empty sse stream".into(), )); } - Ok(StreamableHttpPostResponse::Sse(event_stream, session_id)) } Some(content_type) if content_type.starts_with(JSON_MIME_TYPE) => { @@ -460,7 +460,12 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } } - Ok(sse_stream_from_body(body_stream)) + let maximum_response_bytes = headers + .get(HEADER_MCP_PROTOCOL_VERSION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|version| version == ProtocolVersion::V_2026_07_28.as_str()) + .then_some(MAX_MCP_STDIO_LINE_BYTES); + Ok(sse_stream_from_body(body_stream, maximum_response_bytes)) } } @@ -793,13 +798,124 @@ async fn collect_body( fn sse_stream_from_body( body_stream: HttpResponseBodyStream, + maximum_event_bytes: Option, ) -> BoxStream<'static, std::result::Result> { - SseStream::from_bytes_stream(stream::unfold(body_stream, |mut body_stream| async move { - match body_stream.recv().await { - Ok(Some(bytes)) => Some((Ok(Bytes::from(bytes)), body_stream)), - Ok(None) => None, - Err(error) => Some((Err(io::Error::other(error)), body_stream)), - } - })) + let modern_session = maximum_event_bytes.is_some(); + SseStream::from_bytes_stream(stream::unfold( + (body_stream, SseEventSizeLimit::new(maximum_event_bytes)), + |(mut body_stream, mut size_limit)| async move { + match body_stream.recv().await { + Ok(Some(bytes)) => { + if let Err(error) = size_limit.observe(&bytes) { + Some((Err(error), (body_stream, size_limit))) + } else { + Some((Ok(Bytes::from(bytes)), (body_stream, size_limit))) + } + } + Ok(None) => None, + Err(error) => Some((Err(io::Error::other(error)), (body_stream, size_limit))), + } + }, + )) + .map(move |event| { + event.map(|mut event| { + if let Some(payload) = event.data.as_deref() + && let Some(normalized) = normalize_sse_jsonrpc_message(payload, modern_session) + { + event.data = Some(normalized); + } + event + }) + }) .boxed() } + +struct SseEventSizeLimit { + maximum_bytes: Option, + retained_bytes: usize, + line_bytes: usize, + line_is_comment: bool, + previous_was_carriage_return: bool, + failed: bool, +} + +impl SseEventSizeLimit { + fn new(maximum_bytes: Option) -> Self { + Self { + maximum_bytes, + retained_bytes: 0, + line_bytes: 0, + line_is_comment: false, + previous_was_carriage_return: false, + failed: false, + } + } + + fn observe(&mut self, bytes: &[u8]) -> io::Result<()> { + if self.failed { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "oversized MCP SSE event was already rejected", + )); + } + let Some(maximum_bytes) = self.maximum_bytes else { + return Ok(()); + }; + + for &byte in bytes { + if self.previous_was_carriage_return { + self.previous_was_carriage_return = false; + if byte == b'\n' { + continue; + } + } + + match byte { + b'\r' => { + self.finish_line(maximum_bytes)?; + self.previous_was_carriage_return = true; + } + b'\n' => self.finish_line(maximum_bytes)?, + _ => { + if self.line_bytes == 0 { + self.line_is_comment = byte == b':'; + } + self.line_bytes = self.line_bytes.saturating_add(1); + self.check_limit(maximum_bytes)?; + } + } + } + Ok(()) + } + + fn finish_line(&mut self, maximum_bytes: usize) -> io::Result<()> { + if self.line_bytes == 0 { + self.retained_bytes = 0; + } else if !self.line_is_comment { + // The SSE parser inserts a newline when joining multiple data fields. + self.retained_bytes = self + .retained_bytes + .saturating_add(self.line_bytes) + .saturating_add(1); + } + + self.line_bytes = 0; + self.line_is_comment = false; + self.check_limit(maximum_bytes) + } + + fn check_limit(&mut self, maximum_bytes: usize) -> io::Result<()> { + if self.retained_bytes.saturating_add(self.line_bytes) > maximum_bytes { + self.failed = true; + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("MCP response body exceeds {maximum_bytes} bytes"), + )); + } + Ok(()) + } +} + +#[cfg(test)] +#[path = "http_client_adapter_tests.rs"] +mod tests; diff --git a/codex-rs/rmcp-client/src/http_client_adapter_tests.rs b/codex-rs/rmcp-client/src/http_client_adapter_tests.rs new file mode 100644 index 0000000000..261eff9c9e --- /dev/null +++ b/codex-rs/rmcp-client/src/http_client_adapter_tests.rs @@ -0,0 +1,201 @@ +use std::io::ErrorKind; + +use pretty_assertions::assert_eq; + +use super::SseEventSizeLimit; + +#[test] +fn lf_terminators_reset_the_event_limit() { + let mut limit = SseEventSizeLimit::new(Some(8)); + + limit + .observe(b"data: a\n\ndata: b\n\n") + .expect("LF-terminated events must have independent size limits"); + + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 0)); +} + +#[test] +fn carriage_return_terminators_reset_the_event_limit() { + let mut limit = SseEventSizeLimit::new(Some(8)); + + limit + .observe(b"data: a\r\rdata: b\r\r") + .expect("CR-terminated events must have independent size limits"); + + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 0)); +} + +#[test] +fn crlf_terminators_split_across_chunks_reset_the_event_limit() { + let mut limit = SseEventSizeLimit::new(Some(8)); + + limit + .observe(b"data: a\r") + .expect("a CR must finish the first event field"); + limit + .observe(b"\n\r") + .expect("a split CRLF must not finish another field"); + limit + .observe(b"\ndata: b\r") + .expect("the blank CRLF must reset the first event"); + limit + .observe(b"\n\r\n") + .expect("the second split CRLF event must remain within the limit"); + + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 0)); +} + +#[test] +fn event_at_the_exact_size_limit_is_accepted() { + let mut limit = SseEventSizeLimit::new(Some(9)); + + limit + .observe(b"data: ab\n\n") + .expect("an event at the exact limit must be accepted"); + + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 0)); +} + +#[test] +fn completed_keepalive_comments_do_not_accumulate() { + let mut limit = SseEventSizeLimit::new(Some(6)); + + limit + .observe(&b": ping\n".repeat(/*n*/ 64)) + .expect("completed keepalive comments must not accumulate"); + + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 0)); +} + +#[test] +fn split_keepalive_comments_are_discarded_only_when_complete() { + let mut limit = SseEventSizeLimit::new(Some(6)); + + limit + .observe(b": pi") + .expect("an incomplete comment within the limit must be retained"); + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 4)); + + limit + .observe(b"ng\n: ping\n") + .expect("only completed comments may be excluded from the limit"); + + assert_eq!((limit.retained_bytes, limit.line_bytes), (0, 0)); +} + +#[test] +fn comments_do_not_reset_accumulated_event_data() { + let mut limit = SseEventSizeLimit::new(Some(14)); + + limit + .observe(b"data: a\n") + .expect("the first data field must fit"); + limit + .observe(b": ping\n") + .expect("a completed comment must not count as event data"); + + let error = limit + .observe(b"data: b\n") + .expect_err("a comment must not reset previously retained data"); + + assert_eq!( + (error.kind(), error.to_string()), + ( + ErrorKind::InvalidData, + "MCP response body exceeds 14 bytes".to_string(), + ) + ); +} + +#[test] +fn an_unterminated_comment_remains_size_limited() { + let mut limit = SseEventSizeLimit::new(Some(6)); + + limit + .observe(b": ping") + .expect("a comment at the exact limit must be accepted"); + + let error = limit + .observe(b"!") + .expect_err("an unterminated comment must not bypass the limit"); + + assert_eq!( + (error.kind(), error.to_string()), + ( + ErrorKind::InvalidData, + "MCP response body exceeds 6 bytes".to_string(), + ) + ); +} + +#[test] +fn multiline_data_counts_parser_inserted_newlines() { + let mut limit = SseEventSizeLimit::new(Some(18)); + + let error = limit + .observe(b"data: aaa\ndata: bbb\n\n") + .expect_err("joined multiline event data must remain size limited"); + + assert_eq!( + (error.kind(), error.to_string()), + ( + ErrorKind::InvalidData, + "MCP response body exceeds 18 bytes".to_string(), + ) + ); +} + +#[test] +fn legacy_sse_streams_remain_unlimited() { + let mut limit = SseEventSizeLimit::new(/*maximum_bytes*/ None); + + limit + .observe(&[b'x'; 128]) + .expect("legacy SSE streams must not gain a message limit"); + + assert_eq!( + ( + limit.retained_bytes, + limit.line_bytes, + limit.line_is_comment, + limit.previous_was_carriage_return, + limit.failed, + ), + (0, 0, false, false, false) + ); +} + +#[test] +fn rejected_events_remain_rejected() { + let mut limit = SseEventSizeLimit::new(Some(8)); + + limit + .observe(b"data: abc") + .expect_err("an oversized event must be rejected"); + + let error = limit + .observe(b"") + .expect_err("a rejected event must not be resumed"); + + assert_eq!( + (error.kind(), error.to_string()), + ( + ErrorKind::InvalidData, + "oversized MCP SSE event was already rejected".to_string(), + ) + ); +} + +#[test] +fn size_accounting_saturates_without_overflow() { + let mut limit = SseEventSizeLimit::new(Some(usize::MAX - 1)); + limit.retained_bytes = usize::MAX - 2; + limit.line_bytes = 1; + + let error = limit + .observe(b"x") + .expect_err("saturating counts must still reject an oversized event"); + + assert_eq!((error.kind(), limit.failed), (ErrorKind::InvalidData, true)); +} diff --git a/codex-rs/rmcp-client/src/incoming_jsonrpc.rs b/codex-rs/rmcp-client/src/incoming_jsonrpc.rs index 6e797bbccc..209d0d8f94 100644 --- a/codex-rs/rmcp-client/src/incoming_jsonrpc.rs +++ b/codex-rs/rmcp-client/src/incoming_jsonrpc.rs @@ -1,7 +1,8 @@ -//! Compatibility decoding for modern discovery responses whose server identity -//! is represented in namespaced metadata instead of a top-level field. +//! Compatibility decoding for server messages whose modern wire shapes are +//! not yet represented correctly by rmcp's untagged result union. use rmcp::model::DiscoverResult; +use rmcp::model::InputRequiredResult; use rmcp::model::JsonRpcResponse; use rmcp::model::ServerJsonRpcMessage; use rmcp::model::ServerResult; @@ -11,38 +12,88 @@ use serde_json::Value; // https://github.com/modelcontextprotocol/rust-sdk/issues/1039 const SERVER_INFO_METADATA_KEY: &str = "io.modelcontextprotocol/serverInfo"; -/// Decode at the transport boundary, before rmcp's untagged result union can -/// mistake a metadata-only discovery response for a completed tool call. +/// Decode at the transport boundary, before rmcp's untagged `ServerResult` +/// deserializer can mistake a modern result for a completed tool call. pub(crate) fn deserialize_incoming_jsonrpc_message( bytes: &[u8], ) -> serde_json::Result { let mut message: Value = serde_json::from_slice(bytes)?; - let Some(result) = message.get_mut("result").and_then(Value::as_object_mut) else { + let changed_discovery = normalize_discovery_server_info(&mut message); + + let Some(result) = message.get("result") else { return serde_json::from_value(message); }; - if !result.contains_key("supportedVersions") || !result.contains_key("capabilities") { - return serde_json::from_value(message); + if changed_discovery { + let response: JsonRpcResponse = serde_json::from_value(message)?; + return Ok(ServerJsonRpcMessage::response( + ServerResult::DiscoverResult(response.result), + response.id, + )); + } + + if result.get("resultType").and_then(Value::as_str) == Some("input_required") { + let response: JsonRpcResponse = serde_json::from_value(message)?; + return Ok(ServerJsonRpcMessage::response( + ServerResult::InputRequiredResult(response.result), + response.id, + )); + } + + serde_json::from_value(message) +} + +/// Normalize an SSE payload before rmcp parses it internally. SSE does not +/// expose a typed-message hook, so response-level MRTR metadata cannot be +/// retained until the upstream untagged-result decoder is fixed. +pub(crate) fn normalize_sse_jsonrpc_message(payload: &str, modern_session: bool) -> Option { + if !modern_session { + return None; + } + + let mut message: Value = serde_json::from_str(payload).ok()?; + let changed_discovery = normalize_discovery_server_info(&mut message); + let changed_mrtr = if message + .pointer("/result/resultType") + .and_then(Value::as_str) + == Some("input_required") + { + message + .get_mut("result") + .and_then(Value::as_object_mut) + .is_some_and(|result| result.remove("_meta").is_some()) + } else { + false + }; + + if changed_discovery || changed_mrtr { + serde_json::to_string(&message).ok() + } else { + None + } +} + +fn normalize_discovery_server_info(message: &mut Value) -> bool { + let Some(result) = message.get_mut("result").and_then(Value::as_object_mut) else { + return false; + }; + if result.contains_key("serverInfo") + || !result.contains_key("supportedVersions") + || !result.contains_key("capabilities") + { + return false; } let Some(server_info) = result .get("_meta") .and_then(Value::as_object) - .and_then(|metadata| metadata.get(SERVER_INFO_METADATA_KEY)) + .and_then(|meta| meta.get(SERVER_INFO_METADATA_KEY)) .cloned() else { - return serde_json::from_value(message); + return false; }; - - if !result.contains_key("serverInfo") { - result.insert("serverInfo".to_owned(), server_info); - } - - let response: JsonRpcResponse = serde_json::from_value(message)?; - Ok(ServerJsonRpcMessage::response( - ServerResult::DiscoverResult(response.result), - response.id, - )) + result.insert("serverInfo".to_owned(), server_info); + true } #[cfg(test)] diff --git a/codex-rs/rmcp-client/src/incoming_jsonrpc_tests.rs b/codex-rs/rmcp-client/src/incoming_jsonrpc_tests.rs index 04e42ae2ca..1750ccaca9 100644 --- a/codex-rs/rmcp-client/src/incoming_jsonrpc_tests.rs +++ b/codex-rs/rmcp-client/src/incoming_jsonrpc_tests.rs @@ -4,6 +4,7 @@ use rmcp::model::ServerResult; use serde_json::json; use super::deserialize_incoming_jsonrpc_message; +use super::normalize_sse_jsonrpc_message; #[test] fn discovery_accepts_metadata_namespaced_server_identity() { @@ -63,3 +64,73 @@ fn extension_result_with_discovery_like_fields_remains_a_custom_result() { }; assert_eq!(custom.0, result); } + +#[test] +fn input_required_discriminator_wins_over_tool_result_metadata() { + let message = json!({ + "jsonrpc": "2.0", + "id": "tool-call", + "result": { + "resultType": "input_required", + "requestState": "opaque-server-state", + "_meta": {"retained": "response-metadata"}, + }, + }); + + let decoded = deserialize_incoming_jsonrpc_message(&serde_json::to_vec(&message).unwrap()) + .expect("input-required result must decode"); + let JsonRpcMessage::Response(response) = decoded else { + panic!("expected an input-required response"); + }; + let ServerResult::InputRequiredResult(result) = response.result else { + panic!("input_required must not become an empty completed tool result"); + }; + assert_eq!(result.request_state.as_deref(), Some("opaque-server-state")); + assert_eq!( + result.meta.and_then(|meta| meta.get("retained").cloned()), + Some(json!("response-metadata")) + ); +} + +#[test] +fn sse_discovery_promotes_metadata_namespaced_identity() { + let message = json!({ + "jsonrpc": "2.0", + "id": 1, + "result": { + "resultType": "complete", + "supportedVersions": ["2026-07-28"], + "capabilities": {}, + "ttlMs": 0, + "cacheScope": "private", + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "sse-server", + "version": "1.0.0", + }, + }, + }, + }); + let normalized = + normalize_sse_jsonrpc_message(&message.to_string(), /*modern_session*/ true) + .expect("metadata-only SSE discovery must be normalized"); + let normalized: serde_json::Value = serde_json::from_str(&normalized).unwrap(); + assert_eq!(normalized["result"]["serverInfo"]["name"], "sse-server"); +} + +#[test] +fn legacy_sse_payload_preserves_intermediate_response_metadata() { + let message = json!({ + "jsonrpc": "2.0", + "id": "legacy-extension", + "result": { + "resultType": "input_required", + "_meta": {"preserved": "legacy-extension-metadata"}, + }, + }); + + assert_eq!( + normalize_sse_jsonrpc_message(&message.to_string(), /*modern_session*/ false), + None + ); +} diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index bf13566ea2..a1e14845f2 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -38,6 +38,7 @@ use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::ListToolsResult; use rmcp::model::PaginatedRequestParams; +use rmcp::model::ProtocolVersion; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; @@ -294,7 +295,7 @@ impl From for ElicitationResponse { Self { action: value.action, content: value.content, - meta: None, + meta: value.meta.map(|meta| Value::Object(meta.0)), } } } @@ -303,6 +304,10 @@ impl From for ElicitResult { fn from(value: ElicitationResponse) -> Self { let mut result = Self::new(value.action); result.content = value.content; + result.meta = value.meta.and_then(|meta| match meta { + Value::Object(meta) => Some(rmcp::model::MetaObject::from(meta)), + _ => None, + }); result } } @@ -670,10 +675,22 @@ impl RmcpClient { timeout: Option, ) -> Result { self.refresh_oauth_if_needed().await?; + let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); - async move { service.peer().read_resource(params).await }.boxed() + async move { + let modern_session = requested_modern + && service.peer().peer_info().is_some_and(|info| { + info.protocol_version == ProtocolVersion::V_2026_07_28 + }); + if modern_session { + service.read_resource(params).await + } else { + service.peer().read_resource(params).await + } + } + .boxed() }) .await?; self.persist_oauth_tokens().await; @@ -708,11 +725,20 @@ impl RmcpClient { }; let mut rmcp_params = CallToolRequestParams::new(name); rmcp_params.arguments = arguments; + let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; let result = self .run_service_operation("tools/call", timeout, move |service| { - let rmcp_params = rmcp_params.clone(); + let mut rmcp_params = rmcp_params.clone(); let meta = meta.clone(); async move { + let modern_session = requested_modern + && service.peer().peer_info().is_some_and(|info| { + info.protocol_version == ProtocolVersion::V_2026_07_28 + }); + if modern_session { + rmcp_params.meta = meta; + return service.call_tool(rmcp_params).await; + } let mut options = rmcp::service::PeerRequestOptions::no_options(); options.meta = meta; let result = service @@ -1281,6 +1307,10 @@ impl RmcpClient { initialize_context.timeout, ) .await?; + service + .peer() + .peer_info() + .ok_or_else(|| anyhow!("recovered handshake succeeded but server info was missing"))?; { let mut guard = self.state.lock().await; diff --git a/codex-rs/rmcp-client/tests/mcp_2026_discovery.rs b/codex-rs/rmcp-client/tests/mcp_2026_discovery.rs index de6e63fdfe..ccf9f2c8bd 100644 --- a/codex-rs/rmcp-client/tests/mcp_2026_discovery.rs +++ b/codex-rs/rmcp-client/tests/mcp_2026_discovery.rs @@ -304,6 +304,63 @@ async fn modern_discovery_accepts_metadata_namespaced_server_identity() -> anyho Ok(()) } +#[tokio::test] +async fn modern_discovery_accepts_native_sse_responses() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + assert_eq!(body["method"], "server/discover"); + let response = modern_discover_result(&body); + ResponseTemplate::new(200).set_body_raw( + format!(": keepalive\n\nevent: message\ndata: {response}\n\n"), + "text/event-stream; charset=utf-8", + ) + }) + .expect(1) + .mount(&server) + .await; + + let client = create_client(&server, McpProtocolMode::V20260728).await?; + initialize_client(&client).await?; + client.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn modern_discovery_accepts_metadata_namespaced_server_identity_over_sse() +-> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(|request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + assert_eq!(body["method"], "server/discover"); + let mut response = modern_discover_result(&body); + let server_info = response["result"] + .as_object_mut() + .expect("discovery object") + .remove("serverInfo") + .expect("top-level identity"); + response["result"]["_meta"] = json!({ + "io.modelcontextprotocol/serverInfo": server_info, + }); + ResponseTemplate::new(200).set_body_raw( + format!("event: message\ndata: {response}\n\n"), + "text/event-stream; charset=utf-8", + ) + }) + .expect(1) + .mount(&server) + .await; + + let client = create_client(&server, McpProtocolMode::V20260728).await?; + initialize_client(&client).await?; + client.shutdown().await; + Ok(()) +} + #[tokio::test] async fn modern_http_rejects_oversized_discovery_and_tool_json_responses() -> anyhow::Result<()> { for oversized_method in ["server/discover", "tools/list", "tools/list:error"] { diff --git a/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs b/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs new file mode 100644 index 0000000000..785aa4e195 --- /dev/null +++ b/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs @@ -0,0 +1,292 @@ +use std::collections::HashMap; +use std::ffi::OsString; +use std::sync::Arc; +use std::time::Duration; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::ExecutorStdioServerLauncher; +use codex_rmcp_client::LocalStdioServerLauncher; +use codex_rmcp_client::McpProtocolMode; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StdioServerLauncher; +use futures::FutureExt; +use rmcp::model::ClientCapabilities; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::ProtocolVersion; +use serde_json::Value; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const MAX_MCP_MESSAGE_BYTES: usize = 8 * 1024 * 1024; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum OversizedHttpResponse { + DiscoveryJson, + ToolJson, + ToolError, + ToolSseEvent, +} + +fn initialize_params() -> InitializeRequestParams { + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("codex-message-limit-test", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18) +} + +async fn initialize(client: &RmcpClient) -> anyhow::Result<()> { + client + .initialize( + initialize_params(), + Some(Duration::from_secs(15)), + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({})), + meta: None, + }) + } + .boxed() + }), + ) + .await?; + Ok(()) +} + +fn discovery_response(request: &Value, padding: Option<&str>) -> Value { + json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": { + "resultType": "complete", + "supportedVersions": ["2026-07-28"], + "capabilities": {"tools": {}}, + "instructions": padding.unwrap_or_default(), + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "message-limit-test", + "version": "1.0.0", + }, + }, + "ttlMs": 0, + "cacheScope": "private", + }, + }) +} + +fn tools_response(request: &Value, padding: &str) -> Value { + json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": { + "resultType": "complete", + "tools": [{ + "name": "oversized_tool", + "description": padding, + "inputSchema": {"type": "object"}, + }], + }, + }) +} + +async fn http_client(server: &MockServer, mode: McpProtocolMode) -> anyhow::Result { + RmcpClient::new_streamable_http_client_with_protocol_mode( + "message-limits", + &format!("{}/mcp", server.uri()), + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + mode, + ) + .await +} + +#[tokio::test] +async fn modern_http_rejects_oversized_json_error_and_sse_bodies() -> anyhow::Result<()> { + for oversized in [ + OversizedHttpResponse::DiscoveryJson, + OversizedHttpResponse::ToolJson, + OversizedHttpResponse::ToolError, + OversizedHttpResponse::ToolSseEvent, + ] { + let server = MockServer::start().await; + let padding = Arc::new("x".repeat(MAX_MCP_MESSAGE_BYTES + 1)); + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("JSON-RPC request"); + match body["method"].as_str() { + Some("server/discover") => { + ResponseTemplate::new(200).set_body_json(discovery_response( + &body, + (oversized == OversizedHttpResponse::DiscoveryJson) + .then_some(padding.as_str()), + )) + } + Some("tools/list") if oversized == OversizedHttpResponse::ToolError => { + ResponseTemplate::new(500).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "error": {"code": -32603, "message": padding.as_str()}, + })) + } + Some("tools/list") if oversized == OversizedHttpResponse::ToolSseEvent => { + ResponseTemplate::new(200).set_body_raw( + format!( + "event: message\ndata: {}\n\n", + tools_response(&body, &padding), + ), + "text/event-stream; charset=utf-8", + ) + } + Some("tools/list") => { + ResponseTemplate::new(200).set_body_json(tools_response(&body, &padding)) + } + other => panic!("unexpected MCP method {other:?}"), + } + }) + .mount(&server) + .await; + + let client = http_client(&server, McpProtocolMode::V20260728).await?; + if oversized == OversizedHttpResponse::DiscoveryJson { + let error = initialize(&client) + .await + .expect_err("oversized discovery must be rejected"); + assert!(format!("{error:#}").contains("8388608 bytes")); + } else { + initialize(&client).await?; + let error = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await + .expect_err("oversized tools response must be rejected"); + let error = format!("{error:#}"); + assert!( + error.contains("8388608 bytes") + || (oversized == OversizedHttpResponse::ToolSseEvent + && (error.contains("timed out") || error.contains("Transport closed"))), + "unexpected {oversized:?} error: {error}" + ); + } + client.shutdown().await; + } + Ok(()) +} + +#[tokio::test] +async fn legacy_http_keeps_existing_large_json_response_behavior() -> anyhow::Result<()> { + let server = MockServer::start().await; + let padding = Arc::new("x".repeat(MAX_MCP_MESSAGE_BYTES + 1)); + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("JSON-RPC request"); + match body["method"].as_str() { + Some("initialize") => ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "result": { + "protocolVersion": "2025-06-18", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "legacy-large-response", "version": "1.0.0"}, + }, + })), + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/list") => { + ResponseTemplate::new(200).set_body_json(tools_response(&body, &padding)) + } + other => panic!("unexpected legacy MCP method {other:?}"), + } + }) + .mount(&server) + .await; + + let client = http_client(&server, McpProtocolMode::Legacy).await?; + initialize(&client).await?; + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(15))) + .await?; + assert_eq!( + tools.tools[0].description.as_deref().map(str::len), + Some(MAX_MCP_MESSAGE_BYTES + 1) + ); + client.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn modern_local_and_executor_stdio_reject_oversized_lines() -> anyhow::Result<()> { + let server = codex_utils_cargo_bin::cargo_bin("test_mcp_2026_stdio_server")?; + + for (modern, executor) in [(false, false), (false, true), (true, false), (true, true)] { + let mode = if modern { + McpProtocolMode::V20260728 + } else { + McpProtocolMode::Legacy + }; + let env = modern.then(|| { + HashMap::from([( + OsString::from("CODEX_MCP_PROTOCOL_VERSION"), + OsString::from("2026-07-28"), + )]) + }); + let launcher: Arc = if executor { + Arc::new(ExecutorStdioServerLauncher::new( + Environment::default_for_tests().get_exec_backend(), + )) + } else { + Arc::new(LocalStdioServerLauncher::new(std::env::current_dir()?)) + }; + let fixture_mode = if modern { + "oversized-stdout" + } else { + "oversized-stdout-legacy" + }; + + let client = RmcpClient::new_stdio_client_with_protocol_mode( + server.clone().into(), + vec![OsString::from(fixture_mode)], + env, + &[], + Some(std::env::current_dir()?.to_string_lossy().into_owned()), + launcher, + mode, + ) + .await?; + initialize(&client).await?; + let result = client + .list_tools(/*params*/ None, Some(Duration::from_secs(10))) + .await; + if !modern && !executor { + let tools = result?; + assert_eq!( + tools.tools[0].description.as_deref().map(str::len), + Some(MAX_MCP_MESSAGE_BYTES + 1), + "legacy local stdio must preserve the existing unbounded native codec" + ); + } else { + assert!( + result.is_err(), + "oversized stdio line must be rejected (modern={modern}, executor={executor})" + ); + } + client.shutdown().await; + } + Ok(()) +} diff --git a/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs b/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs new file mode 100644 index 0000000000..22f55cd34e --- /dev/null +++ b/codex-rs/rmcp-client/tests/mcp_2026_mrtr.rs @@ -0,0 +1,683 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_exec_server::Environment; +use codex_rmcp_client::Elicitation; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::McpProtocolMode; +use codex_rmcp_client::RmcpClient; +use futures::FutureExt; +use pretty_assertions::assert_eq; +use rmcp::model::ClientCapabilities; +use rmcp::model::ElicitRequestParams; +use rmcp::model::ElicitationCapability; +use rmcp::model::FormElicitationCapability; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::ProtocolVersion; +use rmcp::model::ReadResourceRequestParams; +use rmcp::model::UrlElicitationCapability; +use serde_json::Value; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const MODERN_VERSION: &str = "2026-07-28"; +const OPAQUE_STATE: &str = " opaque/\u{2603}/=?base64?literal?=\n"; + +fn discover_response(body: &Value) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "result": { + "resultType": "complete", + "supportedVersions": [MODERN_VERSION], + "capabilities": {"tools": {}, "resources": {}}, + "serverInfo": {"name": "mrtr-test", "version": "1.0.0"}, + "ttlMs": 0, + "cacheScope": "private", + }, + })) +} + +fn result_response(body: &Value, result: Value) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "result": result, + })) +} + +fn sse_result_response(body: &Value, result: Value) -> ResponseTemplate { + let message = json!({ + "jsonrpc": "2.0", + "id": body["id"], + "result": result, + }); + ResponseTemplate::new(200).set_body_raw( + format!("event: message\ndata: {message}\n\n"), + "text/event-stream; charset=utf-8", + ) +} + +fn elicitation_request(mode: &str) -> Value { + let params = match mode { + "form" => json!({ + "mode": "form", + "message": "Confirm the MCP request.", + "requestedSchema": { + "type": "object", + "properties": {"confirmed": {"type": "boolean"}}, + "required": ["confirmed"], + }, + "_meta": {"inputContext": "server-context"}, + }), + "url" => json!({ + "mode": "url", + "message": "Approve the MCP request.", + "url": "https://example.test/approve", + "elicitationId": "approval-123", + "_meta": {"inputContext": "server-context"}, + }), + other => panic!("unexpected elicitation mode: {other}"), + }; + json!({"method": "elicitation/create", "params": params}) +} + +async fn create_client( + server: &MockServer, + elicitation_modes: Arc>>, +) -> anyhow::Result { + let client = RmcpClient::new_streamable_http_client_with_protocol_mode( + "mrtr-test", + &format!("{}/mcp", server.uri()), + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + McpProtocolMode::V20260728, + ) + .await?; + + let mut capabilities = ClientCapabilities::default(); + capabilities.elicitation = Some( + ElicitationCapability::new() + .with_form(FormElicitationCapability::new()) + .with_url(UrlElicitationCapability::new()), + ); + client + .initialize( + InitializeRequestParams::new( + capabilities, + Implementation::new("codex-mrtr-test", "0.0.0"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(5)), + Box::new(move |request_id, request| { + let elicitation_modes = Arc::clone(&elicitation_modes); + async move { + assert_eq!(request_id.to_string(), "confirmation"); + let Elicitation::Mcp(request) = request else { + anyhow::bail!("MRTR input must be a standard MCP elicitation"); + }; + let mode = match request { + ElicitRequestParams::FormElicitationParams { meta, .. } => { + assert_eq!( + meta.and_then(|meta| meta.get("inputContext").cloned()), + Some(json!("server-context")) + ); + "form" + } + ElicitRequestParams::UrlElicitationParams { meta, .. } => { + assert_eq!( + meta.and_then(|meta| meta.get("inputContext").cloned()), + Some(json!("server-context")) + ); + "url" + } + _ => anyhow::bail!("unsupported elicitation mode"), + }; + elicitation_modes + .lock() + .map_err(|_| anyhow::anyhow!("elicitation lock was poisoned"))? + .push(mode.to_owned()); + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(json!({"confirmed": true})), + meta: Some(json!({"clientContext": "preserved"})), + }) + } + .boxed() + }), + ) + .await?; + Ok(client) +} + +#[tokio::test] +async fn modern_tool_mrtr_drives_form_and_url_elicitation_and_preserves_metadata() +-> anyhow::Result<()> { + for mode in ["form", "url"] { + let server = MockServer::start().await; + let calls = Arc::new(Mutex::new(Vec::::new())); + let recorded = Arc::clone(&calls); + + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") => { + recorded.lock().expect("requests lock").push(body.clone()); + assert_eq!( + body.pointer("/params/_meta/requestContext"), + Some(&json!("caller-context")) + ); + assert_eq!( + body.pointer("/params/_meta/io.modelcontextprotocol~1protocolVersion"), + Some(&json!(MODERN_VERSION)) + ); + if body.pointer("/params/inputResponses").is_none() { + result_response( + &body, + json!({ + "resultType": "input_required", + "inputRequests": { + "confirmation": elicitation_request(mode), + }, + "requestState": OPAQUE_STATE, + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "mrtr-test", + "version": "1.0.0", + }, + "responseContext": "preserved", + }, + }), + ) + } else { + assert_eq!( + body.pointer("/params/requestState"), + Some(&json!(OPAQUE_STATE)) + ); + assert_eq!( + body.pointer("/params/inputResponses/confirmation"), + Some(&json!({ + "action": "accept", + "content": {"confirmed": true}, + "_meta": {"clientContext": "preserved"}, + })) + ); + result_response( + &body, + json!({ + "resultType": "complete", + "content": [{"type": "text", "text": "MRTR completed"}], + }), + ) + } + } + other => panic!("unexpected MRTR request: {other:?}"), + } + }) + .expect(3) + .mount(&server) + .await; + + let elicitation_modes = Arc::new(Mutex::new(Vec::new())); + let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; + let result = client + .call_tool( + "confirm".into(), + Some(json!({"mode": mode})), + Some(json!({"requestContext": "caller-context"})), + Some(Duration::from_secs(5)), + ) + .await?; + assert_eq!( + result.content[0].as_text().map(|text| text.text.as_str()), + Some("MRTR completed") + ); + assert_eq!( + *elicitation_modes.lock().expect("elicitation lock"), + vec![mode] + ); + + { + let calls = calls.lock().expect("requests lock"); + assert_eq!(calls.len(), 2); + assert_ne!(calls[0]["id"], calls[1]["id"]); + } + client.shutdown().await; + } + Ok(()) +} + +#[tokio::test] +async fn modern_tool_mrtr_uses_recovered_protocol_after_legacy_session_expiry() -> anyhow::Result<()> +{ + let server = MockServer::start().await; + let calls = Arc::new(Mutex::new(Vec::::new())); + let sessions = Arc::new(Mutex::new(Vec::::new())); + let recorded_calls = Arc::clone(&calls); + let recorded_sessions = Arc::clone(&sessions); + + Mock::given(method("GET")) + .and(path("/mcp")) + .respond_with(ResponseTemplate::new(405)) + .mount(&server) + .await; + + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body["method"].as_str() { + Some("server/discover") => { + let mut sessions = recorded_sessions.lock().expect("sessions lock"); + if sessions.is_empty() { + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "error": {"code": -32601, "message": "method not found"}, + })) + } else { + sessions.push("modern-session".to_owned()); + discover_response(&body).insert_header("mcp-session-id", "modern-session") + } + } + Some("initialize") => { + assert_eq!( + body.pointer("/params/protocolVersion"), + Some(&json!(ProtocolVersion::V_2025_06_18.as_str())) + ); + recorded_sessions + .lock() + .expect("sessions lock") + .push("legacy-session".to_owned()); + ResponseTemplate::new(200) + .set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "result": { + "protocolVersion": ProtocolVersion::V_2025_06_18.as_str(), + "capabilities": {"tools": {}}, + "serverInfo": {"name": "mrtr-test", "version": "1.0.0"}, + }, + })) + .insert_header("mcp-session-id", "legacy-session") + } + Some("notifications/initialized") => ResponseTemplate::new(202), + Some("tools/call") => { + assert_eq!( + body.pointer("/params/_meta/requestContext"), + Some(&json!("caller-context")) + ); + let attempt = { + let mut calls = recorded_calls.lock().expect("requests lock"); + calls.push(body.clone()); + calls.len() + }; + + match attempt { + 1 => { + assert_eq!( + body.pointer( + "/params/_meta/io.modelcontextprotocol~1protocolVersion" + ), + None + ); + ResponseTemplate::new(404) + } + 2 => { + assert_eq!( + body.pointer( + "/params/_meta/io.modelcontextprotocol~1protocolVersion" + ), + Some(&json!(MODERN_VERSION)) + ); + result_response( + &body, + json!({ + "resultType": "input_required", + "inputRequests": { + "confirmation": elicitation_request("form"), + }, + "requestState": OPAQUE_STATE, + }), + ) + } + 3 => { + assert_eq!( + body.pointer( + "/params/_meta/io.modelcontextprotocol~1protocolVersion" + ), + Some(&json!(MODERN_VERSION)) + ); + assert_eq!( + body.pointer("/params/requestState"), + Some(&json!(OPAQUE_STATE)) + ); + assert_eq!( + body.pointer("/params/inputResponses/confirmation"), + Some(&json!({ + "action": "accept", + "content": {"confirmed": true}, + "_meta": {"clientContext": "preserved"}, + })) + ); + result_response( + &body, + json!({ + "resultType": "complete", + "content": [{"type": "text", "text": "recovered MRTR completed"}], + }), + ) + } + other => panic!("unexpected recovered MRTR attempt: {other}"), + } + } + other => panic!("unexpected recovered MRTR request: {other:?}"), + } + }) + .mount(&server) + .await; + + let elicitation_modes = Arc::new(Mutex::new(Vec::new())); + let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; + let result = client + .call_tool( + "confirm".into(), + Some(json!({})), + Some(json!({"requestContext": "caller-context"})), + Some(Duration::from_secs(5)), + ) + .await?; + + assert_eq!( + result.content[0].as_text().map(|text| text.text.as_str()), + Some("recovered MRTR completed") + ); + assert_eq!( + *sessions.lock().expect("sessions lock"), + vec!["legacy-session", "modern-session"] + ); + assert_eq!(calls.lock().expect("requests lock").len(), 3); + assert_eq!( + *elicitation_modes.lock().expect("elicitation lock"), + vec!["form"] + ); + client.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn modern_resource_mrtr_retries_with_opaque_request_state() -> anyhow::Result<()> { + let server = MockServer::start().await; + let calls = Arc::new(Mutex::new(Vec::::new())); + let recorded = Arc::clone(&calls); + + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("resources/read") => { + recorded.lock().expect("requests lock").push(body.clone()); + if body.pointer("/params/inputResponses").is_none() { + result_response( + &body, + json!({ + "resultType": "input_required", + "inputRequests": { + "confirmation": elicitation_request("form"), + }, + "requestState": OPAQUE_STATE, + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "mrtr-test", + "version": "1.0.0", + }, + "responseContext": "preserved", + }, + }), + ) + } else { + assert_eq!( + body.pointer("/params/requestState"), + Some(&json!(OPAQUE_STATE)) + ); + assert_eq!( + body.pointer("/params/inputResponses/confirmation/_meta/clientContext"), + Some(&json!("preserved")) + ); + result_response( + &body, + json!({ + "resultType": "complete", + "contents": [{ + "uri": "memo://requires-confirmation", + "mimeType": "text/plain", + "text": "approved resource", + }], + }), + ) + } + } + other => panic!("unexpected resource MRTR request: {other:?}"), + } + }) + .expect(3) + .mount(&server) + .await; + + let elicitation_modes = Arc::new(Mutex::new(Vec::new())); + let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; + let result = client + .read_resource( + ReadResourceRequestParams::new("memo://requires-confirmation"), + Some(Duration::from_secs(5)), + ) + .await?; + assert_eq!(result.contents.len(), 1); + assert_eq!( + *elicitation_modes.lock().expect("elicitation lock"), + vec!["form"] + ); + + { + let calls = calls.lock().expect("requests lock"); + assert_eq!(calls.len(), 2); + assert_ne!(calls[0]["id"], calls[1]["id"]); + } + client.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn modern_tool_mrtr_accepts_sse_discovery_and_metadata_bearing_rounds() -> anyhow::Result<()> +{ + let server = MockServer::start().await; + let calls = Arc::new(Mutex::new(Vec::::new())); + let recorded = Arc::clone(&calls); + + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body["method"].as_str() { + Some("server/discover") => sse_result_response( + &body, + json!({ + "resultType": "complete", + "supportedVersions": [MODERN_VERSION], + "capabilities": {"tools": {}}, + "_meta": { + "io.modelcontextprotocol/serverInfo": { + "name": "sse-mrtr-test", + "version": "1.0.0", + }, + }, + "ttlMs": 0, + "cacheScope": "private", + }), + ), + Some("tools/call") => { + recorded.lock().expect("requests lock").push(body.clone()); + if body.pointer("/params/inputResponses").is_none() { + sse_result_response( + &body, + json!({ + "resultType": "input_required", + "inputRequests": { + "confirmation": elicitation_request("form"), + }, + "requestState": OPAQUE_STATE, + "_meta": {"responseContext": "sse-round"}, + }), + ) + } else { + assert_eq!( + body.pointer("/params/requestState"), + Some(&json!(OPAQUE_STATE)) + ); + assert_eq!( + body.pointer("/params/inputResponses/confirmation/action"), + Some(&json!("accept")) + ); + sse_result_response( + &body, + json!({ + "resultType": "complete", + "content": [{"type": "text", "text": "SSE MRTR completed"}], + }), + ) + } + } + other => panic!("unexpected SSE MRTR request: {other:?}"), + } + }) + .expect(3) + .mount(&server) + .await; + + let elicitation_modes = Arc::new(Mutex::new(Vec::new())); + let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; + let result = client + .call_tool( + "confirm".to_owned(), + Some(json!({})), + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + + assert_eq!( + result.content[0].as_text().map(|text| text.text.as_str()), + Some("SSE MRTR completed") + ); + assert_eq!( + *elicitation_modes.lock().expect("elicitation lock"), + vec!["form"] + ); + assert_eq!(calls.lock().expect("requests lock").len(), 2); + client.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn modern_tool_mrtr_preserves_state_only_rounds_without_elicitation() -> anyhow::Result<()> { + let server = MockServer::start().await; + let states = Arc::new(Mutex::new(Vec::>::new())); + let recorded = Arc::clone(&states); + + Mock::given(method("POST")) + .and(path("/mcp")) + .respond_with(move |request: &Request| { + let body: Value = request.body_json().expect("valid JSON-RPC request"); + match body["method"].as_str() { + Some("server/discover") => discover_response(&body), + Some("tools/call") => { + let state = body + .pointer("/params/requestState") + .and_then(Value::as_str) + .map(str::to_owned); + recorded.lock().expect("states lock").push(state.clone()); + match state.as_deref() { + None => result_response( + &body, + json!({ + "resultType": "input_required", + "requestState": "state-1", + "_meta": {"responseContext": "first-round"}, + }), + ), + Some("state-1") => result_response( + &body, + json!({ + "resultType": "input_required", + "requestState": "state-2", + "_meta": {"responseContext": "second-round"}, + }), + ), + Some("state-2") => result_response( + &body, + json!({ + "resultType": "complete", + "content": [{"type": "text", "text": "state completed"}], + }), + ), + Some(other) => panic!("unexpected opaque state: {other}"), + } + } + other => panic!("unexpected state-only MRTR request: {other:?}"), + } + }) + .expect(4) + .mount(&server) + .await; + + let elicitation_modes = Arc::new(Mutex::new(Vec::new())); + let client = create_client(&server, Arc::clone(&elicitation_modes)).await?; + let result = client + .call_tool( + "state-only".into(), + Some(json!({})), + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + assert_eq!( + result.content[0].as_text().map(|text| text.text.as_str()), + Some("state completed") + ); + assert!( + elicitation_modes + .lock() + .expect("elicitation lock") + .is_empty() + ); + assert_eq!( + *states.lock().expect("states lock"), + vec![ + None, + Some("state-1".to_string()), + Some("state-2".to_string()) + ] + ); + client.shutdown().await; + Ok(()) +} diff --git a/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs b/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs new file mode 100644 index 0000000000..cc75a2c156 --- /dev/null +++ b/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs @@ -0,0 +1,225 @@ +use std::collections::HashMap; +use std::ffi::OsString; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_exec_server::Environment; +use codex_rmcp_client::Elicitation; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::ExecutorStdioServerLauncher; +use codex_rmcp_client::LocalStdioServerLauncher; +use codex_rmcp_client::McpProtocolMode; +use codex_rmcp_client::RmcpClient; +use codex_rmcp_client::StdioServerLauncher; +use futures::FutureExt; +use pretty_assertions::assert_eq; +use rmcp::model::ClientCapabilities; +use rmcp::model::ElicitRequestParams; +use rmcp::model::ElicitationCapability; +use rmcp::model::FormElicitationCapability; +use rmcp::model::Implementation; +use rmcp::model::InitializeRequestParams; +use rmcp::model::ProtocolVersion; +use serde_json::json; + +async fn exercise_stdio_server( + server_mode: &str, + protocol_mode: McpProtocolMode, + opt_in: bool, + use_executor: bool, +) -> anyhow::Result<()> { + let server = codex_utils_cargo_bin::cargo_bin("test_mcp_2026_stdio_server")?; + let launcher: Arc = if use_executor { + Arc::new(ExecutorStdioServerLauncher::new( + Environment::default_for_tests().get_exec_backend(), + )) + } else { + Arc::new(LocalStdioServerLauncher::new(std::env::current_dir()?)) + }; + let env = opt_in.then(|| { + HashMap::from([( + OsString::from("CODEX_MCP_PROTOCOL_VERSION"), + OsString::from("2026-07-28"), + )]) + }); + let client = RmcpClient::new_stdio_client_with_protocol_mode( + server.into(), + vec![OsString::from(server_mode)], + env, + &[], + /*cwd*/ Some(std::env::current_dir()?.to_string_lossy().into_owned()), + launcher, + protocol_mode, + ) + .await?; + + let mut capabilities = ClientCapabilities::default(); + capabilities.elicitation = + Some(ElicitationCapability::new().with_form(FormElicitationCapability::new())); + let elicitation_count = Arc::new(AtomicUsize::new(0)); + let observed_elicitations = Arc::clone(&elicitation_count); + let legacy_session = server_mode.starts_with("legacy"); + let initialized = client + .initialize( + InitializeRequestParams::new( + capabilities, + Implementation::new("codex-mcp-client", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + Some(Duration::from_secs(5)), + Box::new(move |_request_id, request| { + let observed_elicitations = Arc::clone(&observed_elicitations); + async move { + observed_elicitations.fetch_add(1, Ordering::Relaxed); + let Elicitation::Mcp(ElicitRequestParams::FormElicitationParams { + requested_schema, + .. + }) = request + else { + anyhow::bail!("expected a standard MCP form elicitation"); + }; + let content = if legacy_session { + assert_eq!( + serde_json::to_value(requested_schema)?, + json!({ + "type": "object", + "properties": { + "name": {"type": "string", "default": "John Doe"}, + "age": {"type": "integer", "default": 30}, + "score": {"type": "number", "default": 95.5}, + "status": { + "type": "string", + "enum": ["active", "inactive"], + "default": "active", + }, + "verified": {"type": "boolean", "default": true}, + }, + "required": [], + }), + ); + json!({ + "name": "John Doe", + "age": 30, + "score": 95.5, + "status": "active", + "verified": true, + }) + } else { + json!({"approved": true}) + }; + Ok(ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(content), + meta: None, + }) + } + .boxed() + }), + ) + .await?; + + let expected_version = if legacy_session { + ProtocolVersion::V_2025_06_18 + } else { + ProtocolVersion::V_2026_07_28 + }; + assert_eq!(initialized.protocol_version, expected_version); + assert_eq!( + initialized.server_info.name, + if legacy_session { + "legacy-stdio-test" + } else { + "strict-stdio-test" + } + ); + let tools = client + .list_tools(/*params*/ None, Some(Duration::from_secs(5))) + .await?; + assert_eq!( + tools + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect::>(), + vec!["echo"] + ); + let result = client + .call_tool( + "echo".to_owned(), + Some(json!({"message": "hello stdio"})), + /*meta*/ None, + Some(Duration::from_secs(5)), + ) + .await?; + assert_eq!( + result.content[0].as_text().map(|text| text.text.as_str()), + Some(if legacy_session { + "legacy approved" + } else { + "modern approved" + }) + ); + assert_eq!(elicitation_count.load(Ordering::Relaxed), 1); + client.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn modern_local_stdio_discovers_metadata_only_identity_and_drives_mrtr() -> anyhow::Result<()> +{ + exercise_stdio_server( + "modern", + McpProtocolMode::V20260728, + /*opt_in*/ true, + /*use_executor*/ false, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn modern_executor_stdio_discovers_metadata_only_identity_and_drives_mrtr() +-> anyhow::Result<()> { + exercise_stdio_server( + "modern", + McpProtocolMode::V20260728, + /*opt_in*/ true, + /*use_executor*/ true, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn legacy_stdio_supports_sep1034_defaults_without_the_rollout_flag() -> anyhow::Result<()> { + exercise_stdio_server( + "legacy", + McpProtocolMode::Legacy, + /*opt_in*/ false, + /*use_executor*/ false, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rollout_flag_alone_preserves_legacy_stdio_and_sep1034_defaults() -> anyhow::Result<()> { + exercise_stdio_server( + "legacy", + McpProtocolMode::V20260728, + /*opt_in*/ false, + /*use_executor*/ false, + ) + .await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn modern_stdio_safely_falls_back_to_legacy_elicitation() -> anyhow::Result<()> { + exercise_stdio_server( + "legacy-fallback", + McpProtocolMode::V20260728, + /*opt_in*/ true, + /*use_executor*/ false, + ) + .await +}