mirror of
https://github.com/openai/codex.git
synced 2026-11-12 21:20:34 +00:00
Complete MCP 2026 client support (#35725)
## What changed - Drive multi-round `tools/call` and `resources/read` requests through `input_required` responses, preserving opaque request state and elicitation metadata across JSON, SSE, and stdio transports. - Decode modern discovery and elicitation result shapes while retaining legacy initialization, elicitation defaults, and discovery fallback behavior. - Apply the 8 MiB modern-protocol response limit to JSON, SSE events, and stdio messages, and reject executor stream recovery when output sequence gaps would corrupt the protocol stream. ## Testing Added integration coverage for discovery, multi-round requests, legacy fallbacks, message limits, stdio launchers, SSE responses, session recovery, and executor output recovery. GitOrigin-RevId: d26a08883a2825e8ce2408ac498114b33859225e
This commit is contained in:
@@ -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(())
|
||||
}
|
||||
@@ -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<RoleClient> for ElicitationClientService {
|
||||
) -> Result<ClientResult, rmcp::ErrorData> {
|
||||
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<ElicitResult, rmcp::ErrorData> {
|
||||
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<RequestMetaObject>) -> ElicitRequestParams {
|
||||
ElicitRequestParams::FormElicitationParams {
|
||||
meta,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<u64> {
|
||||
watch::channel(0).1
|
||||
}
|
||||
|
||||
fn subscribe_events(&self) -> ExecProcessEventReceiver {
|
||||
ExecProcessEventReceiver::empty()
|
||||
}
|
||||
|
||||
fn read(
|
||||
&self,
|
||||
_after_seq: Option<u64>,
|
||||
_max_bytes: Option<usize>,
|
||||
_wait_ms: Option<u64>,
|
||||
) -> ExecProcessFuture<'_, ReadResponse> {
|
||||
Box::pin(async { Ok(self.response.clone()) })
|
||||
}
|
||||
|
||||
fn write(&self, _chunk: Vec<u8>) -> 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<ProcessOutputChunk>, 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();
|
||||
|
||||
@@ -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<usize>,
|
||||
) -> BoxStream<'static, std::result::Result<Sse, sse_stream::Error>> {
|
||||
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<usize>,
|
||||
retained_bytes: usize,
|
||||
line_bytes: usize,
|
||||
line_is_comment: bool,
|
||||
previous_was_carriage_return: bool,
|
||||
failed: bool,
|
||||
}
|
||||
|
||||
impl SseEventSizeLimit {
|
||||
fn new(maximum_bytes: Option<usize>) -> 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;
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
@@ -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<ServerJsonRpcMessage> {
|
||||
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<DiscoverResult> = 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<InputRequiredResult> = 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<String> {
|
||||
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<DiscoverResult> = serde_json::from_value(message)?;
|
||||
Ok(ServerJsonRpcMessage::response(
|
||||
ServerResult::DiscoverResult(response.result),
|
||||
response.id,
|
||||
))
|
||||
result.insert("serverInfo".to_owned(), server_info);
|
||||
true
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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<ElicitResult> 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<ElicitationResponse> 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<Duration>,
|
||||
) -> Result<ReadResourceResult> {
|
||||
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;
|
||||
|
||||
@@ -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"] {
|
||||
|
||||
@@ -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> {
|
||||
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<dyn StdioServerLauncher> = 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(())
|
||||
}
|
||||
@@ -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<Mutex<Vec<String>>>,
|
||||
) -> anyhow::Result<RmcpClient> {
|
||||
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::<Value>::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::<Value>::new()));
|
||||
let sessions = Arc::new(Mutex::new(Vec::<String>::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::<Value>::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::<Value>::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::<Option<String>>::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(())
|
||||
}
|
||||
@@ -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<dyn StdioServerLauncher> = 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<_>>(),
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user