mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
Add MCP 2026-07-28 discovery support (#35724)
## What changed - Add an opt-in `mcp_2026_07_28` protocol mode while preserving the legacy lifecycle by default. - Negotiate the new protocol over streamable HTTP with `server/discover`, including bounded responses, redirect protection, and fallback only when a response establishes that the endpoint is legacy-only. - Require stdio servers to opt in with `CODEX_MCP_PROTOCOL_VERSION=2026-07-28`, and add a bounded local stdio transport for the modern lifecycle. - Consume paginated tool, resource, and resource-template catalogs in modern mode, reject repeated cursors, and retain discovered server identity. - Reconnect reusable MCP clients when their selected protocol mode changes. ## Testing - Cover HTTP JSON and SSE discovery, legacy fallback and rejection cases, redirects, retries, response limits, and pagination. - Cover local and executor stdio discovery, protocol markers, message limits, and legacy compatibility. GitOrigin-RevId: f6a78816e127d2a482292d63b91c8384f1595903
This commit is contained in:
@@ -14,6 +14,7 @@ use rmcp::model::ResourceTemplate;
|
||||
use tokio::task::JoinSet;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::pagination::collect_paginated;
|
||||
use crate::rmcp_client::ManagedClient;
|
||||
|
||||
/// The ready clients captured for one model step.
|
||||
@@ -90,28 +91,16 @@ impl McpBindingClients {
|
||||
let client = Arc::clone(&managed.client);
|
||||
let timeout = managed.tool_timeout;
|
||||
join_set.spawn(async move {
|
||||
let mut collected = Vec::new();
|
||||
let mut cursor: Option<String> = None;
|
||||
loop {
|
||||
let params = cursor.as_ref().map(|next| {
|
||||
PaginatedRequestParams::default().with_cursor(Some(next.clone()))
|
||||
});
|
||||
let response = match client.list_resources(params, timeout).await {
|
||||
Ok(result) => result,
|
||||
Err(error) => return (server_name, Err(error)),
|
||||
};
|
||||
collected.extend(response.resources);
|
||||
match response.next_cursor {
|
||||
Some(next) if cursor.as_ref() == Some(&next) => {
|
||||
return (
|
||||
server_name,
|
||||
Err(anyhow!("resources/list returned duplicate cursor")),
|
||||
);
|
||||
let resources =
|
||||
collect_paginated("resources/list", /*overall_timeout*/ None, |params| {
|
||||
let client = Arc::clone(&client);
|
||||
async move {
|
||||
let response = client.list_resources(params, timeout).await?;
|
||||
Ok((response.resources, response.next_cursor))
|
||||
}
|
||||
Some(next) => cursor = Some(next),
|
||||
None => return (server_name, Ok(collected)),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
(server_name, resources)
|
||||
});
|
||||
}
|
||||
collect_resource_results(&mut join_set, "resources").await
|
||||
@@ -131,30 +120,19 @@ impl McpBindingClients {
|
||||
let client = Arc::clone(&managed.client);
|
||||
let timeout = managed.tool_timeout;
|
||||
join_set.spawn(async move {
|
||||
let mut collected = Vec::new();
|
||||
let mut cursor: Option<String> = None;
|
||||
loop {
|
||||
let params = cursor.as_ref().map(|next| {
|
||||
PaginatedRequestParams::default().with_cursor(Some(next.clone()))
|
||||
});
|
||||
let response = match client.list_resource_templates(params, timeout).await {
|
||||
Ok(result) => result,
|
||||
Err(error) => return (server_name, Err(error)),
|
||||
};
|
||||
collected.extend(response.resource_templates);
|
||||
match response.next_cursor {
|
||||
Some(next) if cursor.as_ref() == Some(&next) => {
|
||||
return (
|
||||
server_name,
|
||||
Err(anyhow!(
|
||||
"resources/templates/list returned duplicate cursor"
|
||||
)),
|
||||
);
|
||||
let templates = collect_paginated(
|
||||
"resources/templates/list",
|
||||
/*overall_timeout*/ None,
|
||||
|params| {
|
||||
let client = Arc::clone(&client);
|
||||
async move {
|
||||
let response = client.list_resource_templates(params, timeout).await?;
|
||||
Ok((response.resource_templates, response.next_cursor))
|
||||
}
|
||||
Some(next) => cursor = Some(next),
|
||||
None => return (server_name, Ok(collected)),
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
.await;
|
||||
(server_name, templates)
|
||||
});
|
||||
}
|
||||
collect_resource_results(&mut join_set, "resource templates").await
|
||||
|
||||
@@ -187,6 +187,7 @@ impl McpConnectionSet {
|
||||
let codex_home = config.codex_home.clone();
|
||||
let prefix_mcp_tool_names = config.prefix_mcp_tool_names;
|
||||
let non_prefixed_mcp_tool_servers = config.non_prefixed_mcp_tool_servers.clone();
|
||||
let protocol_mode = config.protocol_mode;
|
||||
let client_elicitation_capability = config.client_elicitation_capability.clone();
|
||||
let tool_plugin_provenance = crate::mcp::tool_plugin_provenance(&config);
|
||||
let auth = auth.as_ref();
|
||||
@@ -292,6 +293,26 @@ impl McpConnectionSet {
|
||||
client_elicitation_capability.clone(),
|
||||
supports_openai_form_elicitation,
|
||||
);
|
||||
let expected_protocol_mode = match &configured_config.transport {
|
||||
McpServerTransportConfig::StreamableHttp { .. } => Some(protocol_mode),
|
||||
McpServerTransportConfig::Stdio { .. }
|
||||
if protocol_mode == crate::McpProtocolMode::Legacy =>
|
||||
{
|
||||
Some(crate::McpProtocolMode::Legacy)
|
||||
}
|
||||
McpServerTransportConfig::Stdio { env, .. } => match env
|
||||
.as_ref()
|
||||
.and_then(|variables| variables.get("CODEX_MCP_PROTOCOL_VERSION"))
|
||||
{
|
||||
None => Some(crate::McpProtocolMode::Legacy),
|
||||
Some(version)
|
||||
if version == rmcp::model::ProtocolVersion::V_2026_07_28.as_str() =>
|
||||
{
|
||||
Some(protocol_mode)
|
||||
}
|
||||
Some(_) => None,
|
||||
},
|
||||
};
|
||||
if let Some(previous_view) =
|
||||
reusable_previous.and_then(|previous| previous.servers.get(&server_name))
|
||||
{
|
||||
@@ -299,7 +320,10 @@ impl McpConnectionSet {
|
||||
if connection
|
||||
.reusable_client(&connection_identity)
|
||||
.await
|
||||
.is_some()
|
||||
.is_some_and(|client| {
|
||||
expected_protocol_mode
|
||||
.is_some_and(|expected| client.client.protocol_mode() == expected)
|
||||
})
|
||||
{
|
||||
servers.insert(
|
||||
server_name.clone(),
|
||||
@@ -346,6 +370,7 @@ impl McpConnectionSet {
|
||||
runtime_auth_provider,
|
||||
client_elicitation_capability.clone(),
|
||||
supports_openai_form_elicitation,
|
||||
protocol_mode,
|
||||
);
|
||||
servers.insert(
|
||||
server_name.clone(),
|
||||
|
||||
@@ -226,6 +226,8 @@ struct RefreshTestTransportFactory {
|
||||
tool: Tool,
|
||||
list_started: Option<Arc<Notify>>,
|
||||
release_list: Option<Arc<Notify>>,
|
||||
next_cursor: Option<String>,
|
||||
list_requests: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl ServerHandler for RefreshTestTransportFactory {
|
||||
@@ -238,13 +240,17 @@ impl ServerHandler for RefreshTestTransportFactory {
|
||||
_request: Option<PaginatedRequestParams>,
|
||||
_context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
|
||||
) -> Result<ListToolsResult, McpError> {
|
||||
self.list_requests
|
||||
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
if let Some(list_started) = &self.list_started {
|
||||
list_started.notify_one();
|
||||
}
|
||||
if let Some(release_list) = &self.release_list {
|
||||
release_list.notified().await;
|
||||
}
|
||||
Ok(ListToolsResult::with_all_items(vec![self.tool.clone()]))
|
||||
let mut result = ListToolsResult::with_all_items(vec![self.tool.clone()]);
|
||||
result.next_cursor = self.next_cursor.clone();
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -346,6 +352,48 @@ impl InProcessTransportFactory for DisconnectingToolsTransportFactory {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_tool_catalog_does_not_follow_pagination_cursor() -> anyhow::Result<()> {
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let client = Arc::new(
|
||||
RmcpClient::new_in_process_client(Arc::new(RefreshTestTransportFactory {
|
||||
tool: create_test_tool("legacy", "first-page").tool,
|
||||
list_started: None,
|
||||
release_list: None,
|
||||
next_cursor: Some("next-page".to_string()),
|
||||
list_requests: Arc::clone(&requests),
|
||||
}))
|
||||
.await?,
|
||||
);
|
||||
client
|
||||
.initialize(
|
||||
InitializeRequestParams::new(
|
||||
ClientCapabilities::default(),
|
||||
Implementation::new("codex-test", "0.0.0-test"),
|
||||
)
|
||||
.with_protocol_version(ProtocolVersion::V_2025_06_18),
|
||||
Some(Duration::from_secs(5)),
|
||||
Box::new(|_, _| async { Err(anyhow!("unexpected elicitation")) }.boxed()),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let tools = list_tools_for_client_uncached(
|
||||
"legacy",
|
||||
/*is_codex_apps_mcp_server*/ false,
|
||||
"test",
|
||||
&client,
|
||||
Some(Duration::from_secs(5)),
|
||||
/*server_instructions*/ None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].tool.name.as_ref(), "first-page");
|
||||
assert_eq!(requests.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||
client.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_test_managed_client(tools: Vec<ToolInfo>) -> ManagedClient {
|
||||
ManagedClient {
|
||||
client: Arc::new(
|
||||
@@ -391,6 +439,8 @@ async fn create_test_manager_with_ready_apps_client(
|
||||
tool: tool.tool.clone(),
|
||||
list_started,
|
||||
release_list,
|
||||
next_cursor: None,
|
||||
list_requests: Arc::new(AtomicUsize::new(0)),
|
||||
}))
|
||||
.await?,
|
||||
);
|
||||
@@ -2568,12 +2618,13 @@ fn reusable_server_identity(
|
||||
runtime_context: &McpRuntimeContext,
|
||||
) -> McpServerConnectionIdentity {
|
||||
let server = EffectiveMcpServer::configured(config.clone());
|
||||
let resolved_environment = runtime_context.resolve_server_environment("docs", config);
|
||||
McpServerConnectionIdentity::new(
|
||||
"docs",
|
||||
&server,
|
||||
OAuthCredentialsStoreMode::default(),
|
||||
AuthKeyringBackendKind::default(),
|
||||
&Ok(None),
|
||||
&resolved_environment,
|
||||
runtime_context,
|
||||
/*runtime_auth_provider*/ None,
|
||||
/*auth*/ None,
|
||||
@@ -2796,6 +2847,139 @@ async fn reconciliation_reuses_an_unchanged_ready_server() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reconciliation_reuses_legacy_stdio_server_with_existing_protocol_marker() {
|
||||
let runtime_context = McpRuntimeContext::new(
|
||||
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
PathBuf::from("/tmp"),
|
||||
);
|
||||
let mut config = reusable_server_config("http://127.0.0.1:1");
|
||||
config.transport = McpServerTransportConfig::Stdio {
|
||||
command: "legacy-server".to_string(),
|
||||
args: Vec::new(),
|
||||
env: Some(HashMap::from([(
|
||||
"CODEX_MCP_PROTOCOL_VERSION".to_string(),
|
||||
"1999-01-01".to_string(),
|
||||
)])),
|
||||
env_vars: Vec::new(),
|
||||
cwd: None,
|
||||
};
|
||||
let previous = manager_with_reusable_ready_server(
|
||||
&config,
|
||||
&runtime_context,
|
||||
vec![create_test_tool("docs", "search")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let reconciled = reconcile_reusable_server(&previous, config, runtime_context).await;
|
||||
|
||||
assert!(previous.shares_test_connection_with(&reconciled, "docs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reconciliation_replaces_connection_when_protocol_mode_changes() {
|
||||
let runtime_context = reusable_server_runtime_context();
|
||||
let config = reusable_server_config("http://127.0.0.1:1");
|
||||
let previous = manager_with_reusable_ready_server(
|
||||
&config,
|
||||
&runtime_context,
|
||||
vec![create_test_tool("docs", "search")],
|
||||
)
|
||||
.await;
|
||||
let codex_home = tempdir().expect("tempdir");
|
||||
let mut mcp_config = crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf());
|
||||
mcp_config.protocol_mode = codex_rmcp_client::McpProtocolMode::V20260728;
|
||||
|
||||
let reconciled = McpConnectionSet::new(
|
||||
Some(&previous),
|
||||
McpPublicationGate::already_published(),
|
||||
McpRuntimeInput {
|
||||
config: Arc::new(mcp_config),
|
||||
plugins_available: false,
|
||||
ready_selected_capability_roots: Vec::new(),
|
||||
mcp_servers: HashMap::from([(
|
||||
"docs".to_string(),
|
||||
EffectiveMcpServer::configured(config),
|
||||
)]),
|
||||
submit_id: "refresh".to_string(),
|
||||
tx_event: None,
|
||||
startup_cancellation_token: CancellationToken::new(),
|
||||
runtime_context,
|
||||
codex_apps_tools_cache: ConnectorRuntimeManager::default(),
|
||||
tool_catalog_cache: McpToolCatalogCache::default(),
|
||||
codex_apps_tools_cache_key: ConnectorRuntimeContextKey::personal(
|
||||
/*account_id*/ None, /*chatgpt_user_id*/ None,
|
||||
),
|
||||
supports_openai_form_elicitation: false,
|
||||
auth: None,
|
||||
codex_apps_auth_manager: None,
|
||||
elicitation_reviewer: None,
|
||||
elicitation_lifecycle: None,
|
||||
},
|
||||
ElicitationRequestRouter::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(!previous.shares_test_connection_with(&reconciled, "docs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabled() {
|
||||
let runtime_context = McpRuntimeContext::new(
|
||||
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
PathBuf::from("/tmp"),
|
||||
);
|
||||
let mut config = reusable_server_config("http://127.0.0.1:1");
|
||||
config.transport = McpServerTransportConfig::Stdio {
|
||||
command: "legacy-server".to_string(),
|
||||
args: Vec::new(),
|
||||
env: None,
|
||||
env_vars: Vec::new(),
|
||||
cwd: None,
|
||||
};
|
||||
let previous = manager_with_reusable_ready_server(
|
||||
&config,
|
||||
&runtime_context,
|
||||
vec![create_test_tool("docs", "search")],
|
||||
)
|
||||
.await;
|
||||
let codex_home = tempdir().expect("tempdir");
|
||||
let mut mcp_config = crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf());
|
||||
mcp_config.protocol_mode = codex_rmcp_client::McpProtocolMode::V20260728;
|
||||
|
||||
let reconciled = McpConnectionSet::new(
|
||||
Some(&previous),
|
||||
McpPublicationGate::already_published(),
|
||||
McpRuntimeInput {
|
||||
config: Arc::new(mcp_config),
|
||||
plugins_available: false,
|
||||
ready_selected_capability_roots: Vec::new(),
|
||||
mcp_servers: HashMap::from([(
|
||||
"docs".to_string(),
|
||||
EffectiveMcpServer::configured(config),
|
||||
)]),
|
||||
submit_id: "refresh".to_string(),
|
||||
tx_event: None,
|
||||
startup_cancellation_token: CancellationToken::new(),
|
||||
runtime_context,
|
||||
codex_apps_tools_cache: ConnectorRuntimeManager::default(),
|
||||
tool_catalog_cache: McpToolCatalogCache::default(),
|
||||
codex_apps_tools_cache_key: ConnectorRuntimeContextKey::personal(
|
||||
/*account_id*/ None, /*chatgpt_user_id*/ None,
|
||||
),
|
||||
supports_openai_form_elicitation: false,
|
||||
auth: None,
|
||||
codex_apps_auth_manager: None,
|
||||
elicitation_reviewer: None,
|
||||
elicitation_lifecycle: None,
|
||||
},
|
||||
ElicitationRequestRouter::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(previous.shares_test_connection_with(&reconciled, "docs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reconciliation_updates_elicitation_policy_without_restarting_ready_server() {
|
||||
let runtime_context = reusable_server_runtime_context();
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub use binding::McpBinding;
|
||||
pub use binding::PreparedMcpCall;
|
||||
pub use codex_rmcp_client::McpProtocolMode;
|
||||
pub use connection_manager::tool_is_model_visible;
|
||||
pub use elicitation::ElicitationLifecycle;
|
||||
pub use elicitation::ElicitationReviewRequest;
|
||||
@@ -91,6 +92,7 @@ pub(crate) mod connection_manager;
|
||||
pub(crate) mod elicitation;
|
||||
pub(crate) mod mcp;
|
||||
mod openai_docs_source_attribution;
|
||||
mod pagination;
|
||||
mod plugin_config;
|
||||
mod resource_client;
|
||||
pub(crate) mod rmcp_client;
|
||||
|
||||
@@ -48,6 +48,7 @@ use rmcp::model::ReadResourceResult;
|
||||
use serde_json::Value;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::McpProtocolMode;
|
||||
use crate::ResolvedMcpCatalog;
|
||||
use crate::connection_manager::McpConnectionSet;
|
||||
use crate::runtime::McpPublicationGate;
|
||||
@@ -157,6 +158,8 @@ pub struct McpConfig {
|
||||
pub prefix_mcp_tool_names: bool,
|
||||
/// MCP servers whose model-visible tool namespaces omit the `mcp__` prefix.
|
||||
pub non_prefixed_mcp_tool_servers: Vec<String>,
|
||||
/// Protocol compatibility policy captured when this MCP configuration is created.
|
||||
pub protocol_mode: McpProtocolMode,
|
||||
/// Client-side elicitation capabilities advertised during MCP initialization.
|
||||
pub client_elicitation_capability: ElicitationCapability,
|
||||
/// Resolved MCP registrations keyed by logical server name.
|
||||
|
||||
@@ -37,6 +37,7 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig {
|
||||
apps_enabled: false,
|
||||
prefix_mcp_tool_names: true,
|
||||
non_prefixed_mcp_tool_servers: Vec::new(),
|
||||
protocol_mode: McpProtocolMode::Legacy,
|
||||
client_elicitation_capability: ElicitationCapability::default(),
|
||||
mcp_server_catalog: ResolvedMcpCatalog::default(),
|
||||
connector_snapshot: codex_connectors::ConnectorSnapshot::default(),
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
use std::collections::HashSet;
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Result;
|
||||
use anyhow::anyhow;
|
||||
use rmcp::model::PaginatedRequestParams;
|
||||
|
||||
pub(crate) async fn collect_paginated<T, F, Fut>(
|
||||
method: &str,
|
||||
overall_timeout: Option<Duration>,
|
||||
mut fetch: F,
|
||||
) -> Result<Vec<T>>
|
||||
where
|
||||
F: FnMut(Option<PaginatedRequestParams>) -> Fut,
|
||||
Fut: Future<Output = Result<(Vec<T>, Option<String>)>>,
|
||||
{
|
||||
let collect = async {
|
||||
let mut collected = Vec::new();
|
||||
let mut cursor = None;
|
||||
let mut seen_cursors = HashSet::new();
|
||||
|
||||
loop {
|
||||
let params = cursor.as_ref().map(|next: &String| {
|
||||
PaginatedRequestParams::default().with_cursor(Some(next.clone()))
|
||||
});
|
||||
let (items, next_cursor) = fetch(params).await?;
|
||||
collected.extend(items);
|
||||
|
||||
let Some(next_cursor) = next_cursor else {
|
||||
return Ok(collected);
|
||||
};
|
||||
if !seen_cursors.insert(next_cursor.clone()) {
|
||||
return Err(anyhow!("{method} returned a repeated pagination cursor"));
|
||||
}
|
||||
cursor = Some(next_cursor);
|
||||
}
|
||||
};
|
||||
|
||||
match overall_timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, collect)
|
||||
.await
|
||||
.map_err(|_| anyhow!("{method} pagination timed out after {timeout:?}"))?,
|
||||
None => collect.await,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "pagination_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,107 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::anyhow;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::collect_paginated;
|
||||
|
||||
#[tokio::test]
|
||||
async fn collects_all_pages_including_an_empty_cursor() {
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let observed = Arc::clone(&requests);
|
||||
|
||||
let pages = collect_paginated("tools/list", /*overall_timeout*/ None, move |params| {
|
||||
let observed = Arc::clone(&observed);
|
||||
async move {
|
||||
let cursor = params.and_then(|params| params.cursor);
|
||||
observed.lock().expect("request lock").push(cursor.clone());
|
||||
match cursor.as_deref() {
|
||||
None => Ok((vec!["first"], Some(String::new()))),
|
||||
Some("") => Ok((vec!["second"], Some("last".to_string()))),
|
||||
Some("last") => Ok((vec!["third"], None)),
|
||||
Some(cursor) => Err(anyhow!("unexpected cursor: {cursor}")),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("paginated request succeeds");
|
||||
|
||||
assert_eq!(pages, vec!["first", "second", "third"]);
|
||||
assert_eq!(
|
||||
*requests.lock().expect("request lock"),
|
||||
vec![None, Some(String::new()), Some("last".to_string())]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_nonconsecutive_repeated_cursors() {
|
||||
let error = collect_paginated(
|
||||
"resources/list",
|
||||
/*overall_timeout*/ None,
|
||||
|params| async move {
|
||||
let cursor = params.and_then(|params| params.cursor);
|
||||
let next = match cursor.as_deref() {
|
||||
None => "first",
|
||||
Some("first") => "second",
|
||||
Some("second") => "first",
|
||||
Some(cursor) => return Err(anyhow!("unexpected cursor: {cursor}")),
|
||||
};
|
||||
Ok((Vec::<()>::new(), Some(next.to_string())))
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("a repeated cursor must fail");
|
||||
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"resources/list returned a repeated pagination cursor"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forwards_page_failures() {
|
||||
let error = collect_paginated(
|
||||
"resources/templates/list",
|
||||
/*overall_timeout*/ None,
|
||||
|_params| async { Err::<(Vec<()>, Option<String>), _>(anyhow!("page failed")) },
|
||||
)
|
||||
.await
|
||||
.expect_err("a page error must fail");
|
||||
|
||||
assert_eq!(error.to_string(), "page failed");
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn applies_one_timeout_across_individually_timely_pages() {
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let observed = Arc::clone(&requests);
|
||||
|
||||
let error = collect_paginated("tools/list", Some(Duration::from_secs(5)), move |params| {
|
||||
let observed = Arc::clone(&observed);
|
||||
async move {
|
||||
let cursor = params.and_then(|params| params.cursor);
|
||||
observed.lock().expect("request lock").push(cursor.clone());
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
|
||||
match cursor.as_deref() {
|
||||
None => Ok((vec!["first"], Some("second".to_string()))),
|
||||
Some("second") => Ok((vec!["second"], Some("third".to_string()))),
|
||||
Some("third") => Ok((vec!["third"], None)),
|
||||
Some(cursor) => Err(anyhow!("unexpected cursor: {cursor}")),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect_err("the combined page duration must exceed the shared timeout");
|
||||
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"tools/list pagination timed out after 5s"
|
||||
);
|
||||
assert_eq!(
|
||||
*requests.lock().expect("request lock"),
|
||||
vec![None, Some("second".to_string()), Some("third".to_string())]
|
||||
);
|
||||
}
|
||||
@@ -26,6 +26,7 @@ use crate::elicitation::ElicitationRequestManager;
|
||||
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
use crate::mcp::ToolPluginProvenance;
|
||||
use crate::openai_docs_source_attribution::maybe_with_openai_docs_source_attribution;
|
||||
use crate::pagination::collect_paginated;
|
||||
use crate::runtime::McpRuntimeContext;
|
||||
use crate::runtime::emit_duration;
|
||||
use crate::server::EffectiveMcpServer;
|
||||
@@ -52,6 +53,7 @@ use codex_protocol::protocol::McpStartupStatus;
|
||||
use codex_protocol::protocol::McpStartupUpdateEvent;
|
||||
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 codex_rmcp_client::ToolWithConnectorId;
|
||||
@@ -280,6 +282,7 @@ struct ManagedClientStartup {
|
||||
runtime_auth_provider: Option<SharedAuthProvider>,
|
||||
client_elicitation_capability: ElicitationCapability,
|
||||
supports_openai_form_elicitation: bool,
|
||||
protocol_mode: McpProtocolMode,
|
||||
cancel_token: CancellationToken,
|
||||
startup_complete: Arc<AtomicBool>,
|
||||
}
|
||||
@@ -300,6 +303,7 @@ impl ManagedClientStartup {
|
||||
runtime_auth_provider,
|
||||
client_elicitation_capability,
|
||||
supports_openai_form_elicitation,
|
||||
protocol_mode,
|
||||
cancel_token,
|
||||
startup_complete,
|
||||
} = self.clone();
|
||||
@@ -329,6 +333,7 @@ impl ManagedClientStartup {
|
||||
runtime_context,
|
||||
resolved_environment,
|
||||
runtime_auth_provider,
|
||||
protocol_mode,
|
||||
),
|
||||
)
|
||||
.await
|
||||
@@ -415,6 +420,7 @@ impl AsyncManagedClient {
|
||||
runtime_auth_provider: Option<SharedAuthProvider>,
|
||||
client_elicitation_capability: ElicitationCapability,
|
||||
supports_openai_form_elicitation: bool,
|
||||
protocol_mode: McpProtocolMode,
|
||||
) -> Self {
|
||||
let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME;
|
||||
let reconnect_server_name = server_name.clone();
|
||||
@@ -441,6 +447,7 @@ impl AsyncManagedClient {
|
||||
runtime_auth_provider,
|
||||
client_elicitation_capability,
|
||||
supports_openai_form_elicitation,
|
||||
protocol_mode,
|
||||
cancel_token: cancel_token.clone(),
|
||||
startup_complete: Arc::clone(&startup_complete),
|
||||
});
|
||||
@@ -584,21 +591,36 @@ pub(crate) async fn list_tools_for_client_uncached(
|
||||
server_instructions: Option<&str>,
|
||||
) -> Result<Vec<ToolInfo>> {
|
||||
let fetch_start = Instant::now();
|
||||
let resp = client
|
||||
.list_tools_with_connector_ids(/*params*/ None, timeout)
|
||||
.await?;
|
||||
let tools = resp
|
||||
.tools
|
||||
.into_iter()
|
||||
.map(|tool| {
|
||||
tool_info_from_listed_tool(
|
||||
server_name,
|
||||
is_codex_apps_mcp_server,
|
||||
server_instructions,
|
||||
tool,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let tools = match client.protocol_mode() {
|
||||
McpProtocolMode::Legacy => {
|
||||
client
|
||||
.list_tools_with_connector_ids(/*params*/ None, timeout)
|
||||
.await?
|
||||
.tools
|
||||
}
|
||||
McpProtocolMode::V20260728 => {
|
||||
collect_paginated("tools/list", timeout, |params| {
|
||||
let client = Arc::clone(client);
|
||||
async move {
|
||||
let response = client
|
||||
.list_tools_with_connector_ids(params, timeout)
|
||||
.await?;
|
||||
Ok((response.tools, response.next_cursor))
|
||||
}
|
||||
})
|
||||
.await?
|
||||
}
|
||||
}
|
||||
.into_iter()
|
||||
.map(|tool| {
|
||||
tool_info_from_listed_tool(
|
||||
server_name,
|
||||
is_codex_apps_mcp_server,
|
||||
server_instructions,
|
||||
tool,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
if is_codex_apps_mcp_server {
|
||||
emit_duration(
|
||||
MCP_TOOLS_FETCH_UNCACHED_DURATION_METRIC,
|
||||
@@ -960,6 +982,7 @@ struct StartServerTaskParams {
|
||||
supports_openai_form_elicitation: bool,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
#[instrument(level = "trace", skip_all, fields(server_name = %server_name))]
|
||||
async fn make_rmcp_client(
|
||||
server_name: &str,
|
||||
@@ -969,6 +992,7 @@ async fn make_rmcp_client(
|
||||
runtime_context: McpRuntimeContext,
|
||||
resolved_environment: std::result::Result<Option<Arc<Environment>>, String>,
|
||||
runtime_auth_provider: Option<SharedAuthProvider>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
) -> Result<RmcpClient, StartupOutcomeError> {
|
||||
let config = server.config().clone();
|
||||
let resolved_environment =
|
||||
@@ -1010,9 +1034,17 @@ async fn make_rmcp_client(
|
||||
};
|
||||
|
||||
let cwd = cwd.map(codex_utils_path_uri::LegacyAppPathString::into_string);
|
||||
RmcpClient::new_stdio_client(command_os, args_os, env_os, &env_vars, cwd, launcher)
|
||||
.await
|
||||
.map_err(|err| StartupOutcomeError::from(anyhow!(err)))
|
||||
RmcpClient::new_stdio_client_with_protocol_mode(
|
||||
command_os,
|
||||
args_os,
|
||||
env_os,
|
||||
&env_vars,
|
||||
cwd,
|
||||
launcher,
|
||||
protocol_mode,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| StartupOutcomeError::from(anyhow!(err)))
|
||||
}
|
||||
McpServerTransportConfig::StreamableHttp {
|
||||
url,
|
||||
@@ -1030,7 +1062,7 @@ async fn make_rmcp_client(
|
||||
Ok(token) => token,
|
||||
Err(error) => return Err(error.into()),
|
||||
};
|
||||
RmcpClient::new_streamable_http_client(
|
||||
RmcpClient::new_streamable_http_client_with_protocol_mode(
|
||||
server_name,
|
||||
&url,
|
||||
resolved_bearer_token,
|
||||
@@ -1040,6 +1072,7 @@ async fn make_rmcp_client(
|
||||
keyring_backend_kind,
|
||||
http_client,
|
||||
runtime_auth_provider,
|
||||
protocol_mode,
|
||||
)
|
||||
.await
|
||||
.map_err(StartupOutcomeError::from)
|
||||
|
||||
@@ -6396,6 +6396,30 @@ async fn to_mcp_config_flows_mcp_tool_prefix_from_feature() -> std::io::Result<(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn to_mcp_config_flows_mcp_2026_feature_from_config() -> std::io::Result<()> {
|
||||
let codex_home = TempDir::new()?;
|
||||
let mut config = Config::load_from_base_config_with_overrides(
|
||||
ConfigToml::default(),
|
||||
ConfigOverrides::default(),
|
||||
codex_home.abs(),
|
||||
)
|
||||
.await?;
|
||||
let plugins_manager = PluginsManager::new(codex_home.path().to_path_buf());
|
||||
|
||||
let mcp_config = config.to_mcp_config(&plugins_manager).await;
|
||||
assert_eq!(mcp_config.protocol_mode, codex_mcp::McpProtocolMode::Legacy);
|
||||
|
||||
let _ = config.features.enable(Feature::Mcp20260728);
|
||||
let mcp_config = config.to_mcp_config(&plugins_manager).await;
|
||||
assert_eq!(
|
||||
mcp_config.protocol_mode,
|
||||
codex_mcp::McpProtocolMode::V20260728
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn to_mcp_config_preserves_auth_elicitation_feature_from_config() -> std::io::Result<()> {
|
||||
let codex_home = TempDir::new()?;
|
||||
|
||||
@@ -78,6 +78,7 @@ use codex_login::AuthManagerConfig;
|
||||
use codex_login::AuthRouteConfig;
|
||||
use codex_mcp::McpConfig;
|
||||
use codex_mcp::McpPluginAttribution;
|
||||
use codex_mcp::McpProtocolMode;
|
||||
use codex_mcp::McpServerRegistration;
|
||||
use codex_mcp::ResolvedMcpCatalog;
|
||||
use codex_memories_read::memory_root;
|
||||
@@ -1734,6 +1735,7 @@ impl Config {
|
||||
} else {
|
||||
Vec::new()
|
||||
},
|
||||
protocol_mode: self.mcp_protocol_mode(),
|
||||
client_elicitation_capability: if self.features.enabled(Feature::AuthElicitation) {
|
||||
ElicitationCapability::new()
|
||||
.with_form(FormElicitationCapability::new())
|
||||
@@ -1756,6 +1758,14 @@ impl Config {
|
||||
|| self.non_prefixed_mcp_tool_servers.is_some()
|
||||
}
|
||||
|
||||
pub fn mcp_protocol_mode(&self) -> McpProtocolMode {
|
||||
if self.features.enabled(Feature::Mcp20260728) {
|
||||
McpProtocolMode::V20260728
|
||||
} else {
|
||||
McpProtocolMode::Legacy
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn rebuild_preserving_session_layers(
|
||||
&self,
|
||||
refreshed_config: &Config,
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
use std::io;
|
||||
use std::io::BufRead;
|
||||
use std::io::Write;
|
||||
|
||||
use anyhow::Result;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
|
||||
fn main() -> Result<()> {
|
||||
assert!(std::env::var_os("CODEX_MCP_PROTOCOL_VERSION").is_none());
|
||||
let mut stdout = io::stdout().lock();
|
||||
|
||||
for line in io::stdin().lock().lines() {
|
||||
let request: Value = serde_json::from_str(&line?)?;
|
||||
let result = match request["method"].as_str() {
|
||||
Some("server/discover") => {
|
||||
assert_eq!(
|
||||
request["params"]["_meta"]["io.modelcontextprotocol/protocolVersion"],
|
||||
"2026-07-28"
|
||||
);
|
||||
json!({
|
||||
"resultType": "complete",
|
||||
"supportedVersions": ["2026-07-28"],
|
||||
"capabilities": {"tools": {}, "resources": {}},
|
||||
"_meta": {
|
||||
"io.modelcontextprotocol/serverInfo": {
|
||||
"name": "strict-stdio-discovery",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
},
|
||||
"ttlMs": 0,
|
||||
"cacheScope": "private",
|
||||
})
|
||||
}
|
||||
Some("tools/list") => json!({
|
||||
"resultType": "complete",
|
||||
"tools": [{
|
||||
"name": "stdio_echo",
|
||||
"inputSchema": {"type": "object"},
|
||||
}],
|
||||
}),
|
||||
Some("resources/list") => json!({
|
||||
"resultType": "complete",
|
||||
"resources": [{
|
||||
"uri": "test://stdio/resource",
|
||||
"name": "stdio resource",
|
||||
}],
|
||||
}),
|
||||
Some("resources/templates/list") => json!({
|
||||
"resultType": "complete",
|
||||
"resourceTemplates": [{
|
||||
"uriTemplate": "test://stdio/{name}",
|
||||
"name": "stdio resource template",
|
||||
}],
|
||||
}),
|
||||
method => anyhow::bail!("unexpected modern stdio discovery method: {method:?}"),
|
||||
};
|
||||
serde_json::to_writer(
|
||||
&mut stdout,
|
||||
&json!({"jsonrpc": "2.0", "id": request["id"], "result": result}),
|
||||
)?;
|
||||
stdout.write_all(b"\n")?;
|
||||
stdout.flush()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -128,6 +128,11 @@ impl TestToolServer {
|
||||
{
|
||||
echo.description = Some(Cow::Owned(format!("Echo from {process_label}.")));
|
||||
}
|
||||
if std::env::var_os("MCP_TEST_OVERSIZED_TOOL_DESCRIPTION").is_some()
|
||||
&& let Some(echo) = tools.iter_mut().find(|tool| tool.name == "echo")
|
||||
{
|
||||
echo.description = Some(Cow::Owned("x".repeat(8 * 1024 * 1024 + 1)));
|
||||
}
|
||||
let resources = vec![Self::memo_resource()];
|
||||
let resource_templates = vec![Self::memo_template()];
|
||||
Self {
|
||||
|
||||
@@ -37,7 +37,6 @@ use rmcp::service::RoleClient;
|
||||
use rmcp::service::RxJsonRpcMessage;
|
||||
use rmcp::service::TxJsonRpcMessage;
|
||||
use rmcp::transport::Transport;
|
||||
use serde_json::from_slice;
|
||||
use serde_json::to_vec;
|
||||
use tokio::runtime::Handle;
|
||||
use tokio::sync::Semaphore;
|
||||
@@ -46,6 +45,8 @@ use tracing::debug;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::incoming_jsonrpc::deserialize_incoming_jsonrpc_message;
|
||||
|
||||
static PROCESS_COUNTER: AtomicUsize = AtomicUsize::new(1);
|
||||
// Tool results can make valid MCP responses large, so keep the protocol
|
||||
// ceiling well above ordinary messages while still bounding hostile input.
|
||||
@@ -426,7 +427,7 @@ impl ExecutorProcessTransport {
|
||||
None => return None,
|
||||
};
|
||||
let line = Self::trim_trailing_carriage_return(line);
|
||||
match from_slice::<RxJsonRpcMessage<RoleClient>>(&line) {
|
||||
match deserialize_incoming_jsonrpc_message(&line) {
|
||||
Ok(message) => return Some(message),
|
||||
Err(error) => {
|
||||
debug!(
|
||||
|
||||
@@ -31,8 +31,14 @@ use reqwest::header::HeaderName;
|
||||
use rmcp::model::ClientJsonRpcMessage;
|
||||
use rmcp::model::ClientNotification;
|
||||
use rmcp::model::ConstString;
|
||||
use rmcp::model::DiscoverRequestMethod;
|
||||
use rmcp::model::ErrorCode;
|
||||
use rmcp::model::ErrorData;
|
||||
use rmcp::model::JsonRpcMessage;
|
||||
use rmcp::model::ProtocolVersion;
|
||||
use rmcp::model::ServerJsonRpcMessage;
|
||||
use rmcp::model::ServerResult;
|
||||
use rmcp::transport::common::http_header::HEADER_MCP_PROTOCOL_VERSION;
|
||||
use rmcp::transport::streamable_http_client::AuthRequiredError;
|
||||
use rmcp::transport::streamable_http_client::InsufficientScopeError;
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpClient;
|
||||
@@ -41,6 +47,9 @@ use rmcp::transport::streamable_http_client::StreamableHttpPostResponse;
|
||||
use sse_stream::Sse;
|
||||
use sse_stream::SseStream;
|
||||
|
||||
use crate::incoming_jsonrpc::deserialize_incoming_jsonrpc_message;
|
||||
use crate::local_stdio_transport::MAX_MCP_STDIO_LINE_BYTES;
|
||||
|
||||
mod www_authenticate;
|
||||
|
||||
use self::www_authenticate::insufficient_scope_challenge;
|
||||
@@ -65,6 +74,8 @@ pub(crate) enum StreamableHttpClientAdapterError {
|
||||
HttpRequest(#[from] ExecServerError),
|
||||
#[error("invalid HTTP header: {0}")]
|
||||
Header(String),
|
||||
#[error("MCP response body exceeds {maximum_bytes} bytes")]
|
||||
ResponseTooLarge { maximum_bytes: usize },
|
||||
}
|
||||
|
||||
impl StreamableHttpClientAdapter {
|
||||
@@ -126,6 +137,18 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
)?;
|
||||
}
|
||||
|
||||
let maximum_response_bytes = (mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE)
|
||||
|| headers
|
||||
.get(HEADER_MCP_PROTOCOL_VERSION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
== Some(ProtocolVersion::V_2026_07_28.as_str()))
|
||||
.then_some(MAX_MCP_STDIO_LINE_BYTES);
|
||||
let redirect_policy = if mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE) {
|
||||
HttpRedirectPolicy::Stop
|
||||
} else {
|
||||
mcp_redirect_policy(&headers)
|
||||
};
|
||||
|
||||
let body = serde_json::to_vec(&message).map_err(StreamableHttpError::Deserialize)?;
|
||||
let has_authorization_header = headers.contains_key(AUTHORIZATION);
|
||||
let response = self
|
||||
@@ -136,7 +159,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
headers: protocol_headers(&headers),
|
||||
body: Some(body.into()),
|
||||
timeout_ms: None,
|
||||
redirect_policy: HttpRedirectPolicy::Follow,
|
||||
redirect_policy,
|
||||
request_id: "buffered-request".to_string(),
|
||||
stream_response: true,
|
||||
})
|
||||
@@ -190,14 +213,42 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
let content_type = response_header(&response.headers, CONTENT_TYPE);
|
||||
let session_id = response_header(&response.headers, HEADER_SESSION_ID);
|
||||
if !status_is_success(response.status) {
|
||||
let body = collect_body(&mut body_stream).await?;
|
||||
let body = collect_body(&mut body_stream, maximum_response_bytes).await?;
|
||||
if !retryable_post_response_status(mcp_method.as_deref(), response.status)
|
||||
&& content_type
|
||||
.as_deref()
|
||||
.is_some_and(|content_type| content_type.starts_with(JSON_MIME_TYPE))
|
||||
&& let Some(message) = parse_json_rpc_error(&body)
|
||||
&& let Some(response_message) = parse_json_rpc_error(&body)
|
||||
{
|
||||
return Ok(StreamableHttpPostResponse::Json(message, session_id));
|
||||
return Ok(StreamableHttpPostResponse::Json(
|
||||
legacy_discovery_fallback_response(
|
||||
&message,
|
||||
response_message,
|
||||
response.status == StatusCode::BAD_REQUEST.as_u16() && !has_session_id,
|
||||
),
|
||||
session_id,
|
||||
));
|
||||
}
|
||||
if mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE)
|
||||
&& !has_session_id
|
||||
&& matches!(
|
||||
StatusCode::from_u16(response.status).ok(),
|
||||
Some(StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED)
|
||||
)
|
||||
&& content_type
|
||||
.as_deref()
|
||||
.is_none_or(|content_type| !content_type.starts_with(JSON_MIME_TYPE))
|
||||
&& let JsonRpcMessage::Request(request) = &message
|
||||
{
|
||||
let legacy_error = ServerJsonRpcMessage::error(
|
||||
ErrorData::new(
|
||||
ErrorCode::METHOD_NOT_FOUND,
|
||||
"legacy MCP endpoint does not support server/discover",
|
||||
None,
|
||||
),
|
||||
Some(request.id.clone()),
|
||||
);
|
||||
return Ok(StreamableHttpPostResponse::Json(legacy_error, session_id));
|
||||
}
|
||||
return Err(StreamableHttpError::UnexpectedServerResponse(
|
||||
format!(
|
||||
@@ -210,17 +261,55 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
}
|
||||
match content_type.as_deref() {
|
||||
Some(content_type) if content_type.starts_with(EVENT_STREAM_MIME_TYPE) => {
|
||||
let event_stream = sse_stream_from_body(body_stream);
|
||||
let mut event_stream = sse_stream_from_body(body_stream);
|
||||
if mcp_method.as_deref() == Some(DiscoverRequestMethod::VALUE) {
|
||||
while let Some(event) = event_stream.next().await {
|
||||
let event = event.map_err(StreamableHttpError::Sse)?;
|
||||
if !matches!(event.event.as_deref(), None | Some("") | Some("message")) {
|
||||
continue;
|
||||
}
|
||||
let Some(data) = event.data.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
if data.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let response = deserialize_incoming_jsonrpc_message(data.as_bytes())
|
||||
.map_err(StreamableHttpError::Deserialize)?;
|
||||
let response = legacy_discovery_fallback_response(
|
||||
&message, response, /*allow_uncorrelated_http_rejection*/ false,
|
||||
);
|
||||
if matches!(
|
||||
&response,
|
||||
JsonRpcMessage::Response(_) | JsonRpcMessage::Error(_)
|
||||
) {
|
||||
return Ok(StreamableHttpPostResponse::Json(response, session_id));
|
||||
}
|
||||
}
|
||||
|
||||
return Err(StreamableHttpError::UnexpectedServerResponse(
|
||||
"empty sse stream".into(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
|
||||
}
|
||||
Some(content_type) if content_type.starts_with(JSON_MIME_TYPE) => {
|
||||
let body = collect_body(&mut body_stream).await?;
|
||||
let message: ServerJsonRpcMessage =
|
||||
serde_json::from_slice(&body).map_err(StreamableHttpError::Deserialize)?;
|
||||
Ok(StreamableHttpPostResponse::Json(message, session_id))
|
||||
let body = collect_body(&mut body_stream, maximum_response_bytes).await?;
|
||||
let response_message = deserialize_incoming_jsonrpc_message(&body)
|
||||
.map_err(StreamableHttpError::Deserialize)?;
|
||||
Ok(StreamableHttpPostResponse::Json(
|
||||
legacy_discovery_fallback_response(
|
||||
&message,
|
||||
response_message,
|
||||
/*allow_uncorrelated_http_rejection*/ false,
|
||||
),
|
||||
session_id,
|
||||
))
|
||||
}
|
||||
_ => {
|
||||
let body = collect_body(&mut body_stream).await?;
|
||||
let body = collect_body(&mut body_stream, maximum_response_bytes).await?;
|
||||
let content_type = content_type.unwrap_or_else(|| "missing-content-type".into());
|
||||
Err(StreamableHttpError::UnexpectedContentType(Some(format!(
|
||||
"{content_type}; body: {}",
|
||||
@@ -254,6 +343,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
session.to_string(),
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
let redirect_policy = mcp_redirect_policy(&headers);
|
||||
|
||||
let response = self
|
||||
.http_client
|
||||
@@ -263,7 +353,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
headers: protocol_headers(&headers),
|
||||
body: None,
|
||||
timeout_ms: None,
|
||||
redirect_policy: HttpRedirectPolicy::Follow,
|
||||
redirect_policy,
|
||||
request_id: "buffered-request".to_string(),
|
||||
stream_response: false,
|
||||
})
|
||||
@@ -326,6 +416,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
StreamableHttpClientAdapterError::Header,
|
||||
)?;
|
||||
}
|
||||
let redirect_policy = mcp_redirect_policy(&headers);
|
||||
|
||||
let (response, body_stream) = self
|
||||
.http_client
|
||||
@@ -335,7 +426,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter {
|
||||
headers: protocol_headers(&headers),
|
||||
body: None,
|
||||
timeout_ms: None,
|
||||
redirect_policy: HttpRedirectPolicy::Follow,
|
||||
redirect_policy,
|
||||
request_id: "buffered-request".to_string(),
|
||||
stream_response: true,
|
||||
})
|
||||
@@ -517,7 +608,12 @@ fn retryable_post_response_status(mcp_method: Option<&str>, status: u16) -> bool
|
||||
is_retryable_http_status(status)
|
||||
&& matches!(
|
||||
mcp_method,
|
||||
Some("initialize" | "notifications/initialized" | "tools/list")
|
||||
Some(
|
||||
DiscoverRequestMethod::VALUE
|
||||
| "initialize"
|
||||
| "notifications/initialized"
|
||||
| "tools/list"
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
@@ -540,8 +636,141 @@ fn parse_json_rpc_error(body: &[u8]) -> Option<ServerJsonRpcMessage> {
|
||||
}
|
||||
}
|
||||
|
||||
fn mcp_redirect_policy(headers: &HeaderMap) -> HttpRedirectPolicy {
|
||||
if headers
|
||||
.get(HEADER_MCP_PROTOCOL_VERSION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
== Some(ProtocolVersion::V_2026_07_28.as_str())
|
||||
{
|
||||
HttpRedirectPolicy::Stop
|
||||
} else {
|
||||
HttpRedirectPolicy::Follow
|
||||
}
|
||||
}
|
||||
|
||||
// rmcp's automatic lifecycle does not yet recognize deployed legacy discovery
|
||||
// rejection shapes. Remove this compatibility shim once the SDK does:
|
||||
// https://github.com/modelcontextprotocol/rust-sdk/issues/1040
|
||||
fn legacy_discovery_fallback_response(
|
||||
request: &ClientJsonRpcMessage,
|
||||
response: ServerJsonRpcMessage,
|
||||
allow_uncorrelated_http_rejection: bool,
|
||||
) -> ServerJsonRpcMessage {
|
||||
let JsonRpcMessage::Request(request) = request else {
|
||||
return response;
|
||||
};
|
||||
if request.request.method() != DiscoverRequestMethod::VALUE {
|
||||
return response;
|
||||
}
|
||||
|
||||
if let JsonRpcMessage::Error(error) = &response
|
||||
&& error.error.code == ErrorCode::METHOD_NOT_FOUND
|
||||
&& error.id.as_ref() != Some(&request.id)
|
||||
{
|
||||
return ServerJsonRpcMessage::error(
|
||||
ErrorData::new(
|
||||
ErrorCode::INVALID_REQUEST,
|
||||
"server/discover method-not-found response did not match its request ID",
|
||||
None,
|
||||
),
|
||||
Some(request.id.clone()),
|
||||
);
|
||||
}
|
||||
|
||||
let only_known_legacy_versions = match &response {
|
||||
JsonRpcMessage::Response(response) if response.id == request.id => match &response.result {
|
||||
ServerResult::DiscoverResult(result) => {
|
||||
only_known_legacy_protocol_versions(&result.supported_versions)
|
||||
}
|
||||
_ => false,
|
||||
},
|
||||
JsonRpcMessage::Error(error) if error.id.as_ref() == Some(&request.id) => {
|
||||
(error.error.code == ErrorCode::UNSUPPORTED_PROTOCOL_VERSION
|
||||
&& error
|
||||
.error
|
||||
.data
|
||||
.as_ref()
|
||||
.and_then(|data| data.get("supported"))
|
||||
.and_then(|supported| {
|
||||
serde_json::from_value::<Vec<ProtocolVersion>>(supported.clone()).ok()
|
||||
})
|
||||
.is_some_and(|supported| only_known_legacy_protocol_versions(&supported)))
|
||||
|| matches!(
|
||||
error.error.code,
|
||||
ErrorCode::UNSUPPORTED_PROTOCOL_VERSION
|
||||
| ErrorCode::INVALID_REQUEST
|
||||
| ErrorCode::INVALID_PARAMS
|
||||
) && explicitly_rejects_modern_protocol_version(&error.error.message)
|
||||
}
|
||||
JsonRpcMessage::Error(error)
|
||||
if error.id.is_none()
|
||||
&& allow_uncorrelated_http_rejection
|
||||
&& error.error.code.0 == -32000 =>
|
||||
{
|
||||
rejects_modern_protocol_at_legacy_http_boundary(&error.error.message)
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if only_known_legacy_versions {
|
||||
ServerJsonRpcMessage::error(
|
||||
ErrorData::new(
|
||||
ErrorCode::METHOD_NOT_FOUND,
|
||||
"MCP discovery only advertised legacy protocol versions",
|
||||
None,
|
||||
),
|
||||
Some(request.id.clone()),
|
||||
)
|
||||
} else {
|
||||
response
|
||||
}
|
||||
}
|
||||
|
||||
fn only_known_legacy_protocol_versions(versions: &[ProtocolVersion]) -> bool {
|
||||
!versions.is_empty()
|
||||
&& versions.iter().all(|version| {
|
||||
ProtocolVersion::KNOWN_VERSIONS.contains(version)
|
||||
&& version < &ProtocolVersion::V_2026_07_28
|
||||
})
|
||||
}
|
||||
|
||||
fn explicitly_rejects_modern_protocol_version(message: &str) -> bool {
|
||||
message
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("unsupported protocol version: 2026-07-28")
|
||||
}
|
||||
|
||||
fn rejects_modern_protocol_at_legacy_http_boundary(message: &str) -> bool {
|
||||
if message == "Bad Request: No valid session ID provided" {
|
||||
return true;
|
||||
}
|
||||
|
||||
let Some(supported) = message
|
||||
.strip_prefix("Bad Request: Unsupported protocol version: 2026-07-28 (supported versions: ")
|
||||
.and_then(|supported| supported.strip_suffix(')'))
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let versions = supported.split(',').map(str::trim).collect::<Vec<_>>();
|
||||
!versions.is_empty()
|
||||
&& versions.contains(&"2025-06-18")
|
||||
&& versions.iter().all(|version| {
|
||||
let bytes = version.as_bytes();
|
||||
bytes.len() == 10
|
||||
&& bytes[4] == b'-'
|
||||
&& bytes[7] == b'-'
|
||||
&& bytes
|
||||
.iter()
|
||||
.enumerate()
|
||||
.all(|(index, byte)| matches!(index, 4 | 7) || byte.is_ascii_digit())
|
||||
&& *version < "2026-07-28"
|
||||
})
|
||||
}
|
||||
|
||||
async fn collect_body(
|
||||
body_stream: &mut HttpResponseBodyStream,
|
||||
maximum_bytes: Option<usize>,
|
||||
) -> std::result::Result<Vec<u8>, StreamableHttpError<StreamableHttpClientAdapterError>> {
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = body_stream
|
||||
@@ -550,6 +779,13 @@ async fn collect_body(
|
||||
.map_err(StreamableHttpClientAdapterError::from)
|
||||
.map_err(StreamableHttpError::Client)?
|
||||
{
|
||||
if let Some(maximum_bytes) = maximum_bytes
|
||||
&& chunk.len() > maximum_bytes.saturating_sub(body.len())
|
||||
{
|
||||
return Err(StreamableHttpError::Client(
|
||||
StreamableHttpClientAdapterError::ResponseTooLarge { maximum_bytes },
|
||||
));
|
||||
}
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok(body)
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
//! Compatibility decoding for modern discovery responses whose server identity
|
||||
//! is represented in namespaced metadata instead of a top-level field.
|
||||
|
||||
use rmcp::model::DiscoverResult;
|
||||
use rmcp::model::JsonRpcResponse;
|
||||
use rmcp::model::ServerJsonRpcMessage;
|
||||
use rmcp::model::ServerResult;
|
||||
use serde_json::Value;
|
||||
|
||||
// Remove this compatibility decoder once rmcp accepts namespaced server identity:
|
||||
// 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.
|
||||
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 {
|
||||
return serde_json::from_value(message);
|
||||
};
|
||||
|
||||
if !result.contains_key("supportedVersions") || !result.contains_key("capabilities") {
|
||||
return serde_json::from_value(message);
|
||||
}
|
||||
|
||||
let Some(server_info) = result
|
||||
.get("_meta")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get(SERVER_INFO_METADATA_KEY))
|
||||
.cloned()
|
||||
else {
|
||||
return serde_json::from_value(message);
|
||||
};
|
||||
|
||||
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,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "incoming_jsonrpc_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,65 @@
|
||||
use pretty_assertions::assert_eq;
|
||||
use rmcp::model::JsonRpcMessage;
|
||||
use rmcp::model::ServerResult;
|
||||
use serde_json::json;
|
||||
|
||||
use super::deserialize_incoming_jsonrpc_message;
|
||||
|
||||
#[test]
|
||||
fn discovery_accepts_metadata_namespaced_server_identity() {
|
||||
let message = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"result": {
|
||||
"resultType": "complete",
|
||||
"supportedVersions": ["2026-07-28"],
|
||||
"capabilities": {"tools": {}},
|
||||
"ttlMs": 0,
|
||||
"cacheScope": "private",
|
||||
"_meta": {
|
||||
"io.modelcontextprotocol/serverInfo": {
|
||||
"name": "modern-server",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
"retained": "metadata",
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
let decoded = deserialize_incoming_jsonrpc_message(&serde_json::to_vec(&message).unwrap())
|
||||
.expect("metadata-only identity must decode as discovery");
|
||||
let JsonRpcMessage::Response(response) = decoded else {
|
||||
panic!("expected a discovery response");
|
||||
};
|
||||
let ServerResult::DiscoverResult(result) = response.result else {
|
||||
panic!("metadata-only identity must not become a completed tool result");
|
||||
};
|
||||
assert_eq!(result.server_info.name, "modern-server");
|
||||
assert_eq!(
|
||||
result.meta.and_then(|meta| meta.get("retained").cloned()),
|
||||
Some(json!("metadata"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extension_result_with_discovery_like_fields_remains_a_custom_result() {
|
||||
let result = json!({
|
||||
"supportedVersions": ["custom-extension-v1"],
|
||||
"capabilities": {"custom": true},
|
||||
});
|
||||
let message = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"result": result,
|
||||
});
|
||||
|
||||
let decoded = deserialize_incoming_jsonrpc_message(&serde_json::to_vec(&message).unwrap())
|
||||
.expect("legacy extension result must deserialize normally");
|
||||
let JsonRpcMessage::Response(response) = decoded else {
|
||||
panic!("expected a legacy extension response");
|
||||
};
|
||||
let ServerResult::CustomResult(custom) = response.result else {
|
||||
panic!("legacy extension response must remain a custom result");
|
||||
};
|
||||
assert_eq!(custom.0, result);
|
||||
}
|
||||
@@ -3,11 +3,14 @@ mod elicitation_client_service;
|
||||
mod executor_process_transport;
|
||||
mod http_client_adapter;
|
||||
mod in_process_transport;
|
||||
mod incoming_jsonrpc;
|
||||
mod local_stdio_transport;
|
||||
mod logging_client_handler;
|
||||
mod oauth;
|
||||
mod oauth_http_client;
|
||||
mod perform_oauth_login;
|
||||
mod program_resolver;
|
||||
mod protocol_mode;
|
||||
mod rmcp_client;
|
||||
mod startup_error;
|
||||
mod stdio_server_launcher;
|
||||
@@ -36,6 +39,7 @@ pub use perform_oauth_login::perform_oauth_login;
|
||||
pub use perform_oauth_login::perform_oauth_login_return_url;
|
||||
pub use perform_oauth_login::perform_oauth_login_return_url_with_http_client;
|
||||
pub use perform_oauth_login::perform_oauth_login_silent;
|
||||
pub use protocol_mode::McpProtocolMode;
|
||||
pub use rmcp::model::ElicitationAction;
|
||||
pub use rmcp_client::Elicitation;
|
||||
pub use rmcp_client::ElicitationResponse;
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
//! Bounded, compatibility-aware transport for modern locally spawned MCP servers.
|
||||
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use memchr::memchr;
|
||||
use rmcp::service::RoleClient;
|
||||
use rmcp::service::RxJsonRpcMessage;
|
||||
use rmcp::service::TxJsonRpcMessage;
|
||||
use rmcp::transport::Transport;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::process::Child;
|
||||
use tokio::process::ChildStderr;
|
||||
use tokio::process::ChildStdin;
|
||||
use tokio::process::ChildStdout;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::debug;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::incoming_jsonrpc::deserialize_incoming_jsonrpc_message;
|
||||
|
||||
/// Match the existing executor stdio transport and the production MCP limit.
|
||||
pub(crate) const MAX_MCP_STDIO_LINE_BYTES: usize = 8 * 1024 * 1024;
|
||||
|
||||
pub(super) struct LocalStdioTransport {
|
||||
child: Child,
|
||||
stdin: Arc<Mutex<Option<ChildStdin>>>,
|
||||
stdout: BufReader<ChildStdout>,
|
||||
pending_line: Vec<u8>,
|
||||
program_name: String,
|
||||
}
|
||||
|
||||
impl LocalStdioTransport {
|
||||
pub(super) fn spawn(
|
||||
mut command: Command,
|
||||
program_name: String,
|
||||
) -> io::Result<(Self, Option<ChildStderr>)> {
|
||||
command
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
let mut child = command.spawn()?;
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| io::Error::other("MCP server stdin was not piped"))?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| io::Error::other("MCP server stdout was not piped"))?;
|
||||
let stderr = child.stderr.take();
|
||||
|
||||
Ok((
|
||||
Self {
|
||||
child,
|
||||
stdin: Arc::new(Mutex::new(Some(stdin))),
|
||||
stdout: BufReader::new(stdout),
|
||||
pending_line: Vec::new(),
|
||||
program_name,
|
||||
},
|
||||
stderr,
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) fn id(&self) -> Option<u32> {
|
||||
self.child.id()
|
||||
}
|
||||
|
||||
async fn receive_message(&mut self) -> Option<RxJsonRpcMessage<RoleClient>> {
|
||||
loop {
|
||||
let bytes = match self.stdout.fill_buf().await {
|
||||
Ok(bytes) => bytes,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
"Failed to read MCP server stdout ({}): {error}",
|
||||
self.program_name
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
if bytes.is_empty() {
|
||||
if self.pending_line.is_empty() {
|
||||
return None;
|
||||
}
|
||||
return self.decode_pending_message();
|
||||
}
|
||||
|
||||
let newline = memchr(b'\n', bytes);
|
||||
let content_len = newline.unwrap_or(bytes.len());
|
||||
if content_len > MAX_MCP_STDIO_LINE_BYTES.saturating_sub(self.pending_line.len()) {
|
||||
warn!(
|
||||
"MCP stdio line exceeds {MAX_MCP_STDIO_LINE_BYTES} bytes ({}); closing transport",
|
||||
self.program_name
|
||||
);
|
||||
self.pending_line.clear();
|
||||
return None;
|
||||
}
|
||||
|
||||
self.pending_line.extend_from_slice(&bytes[..content_len]);
|
||||
let consumed = content_len + usize::from(newline.is_some());
|
||||
self.stdout.consume(consumed);
|
||||
|
||||
if newline.is_some()
|
||||
&& let Some(message) = self.decode_pending_message()
|
||||
{
|
||||
return Some(message);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_pending_message(&mut self) -> Option<RxJsonRpcMessage<RoleClient>> {
|
||||
let line = std::mem::take(&mut self.pending_line);
|
||||
let line = line.strip_suffix(b"\r").unwrap_or(&line);
|
||||
match deserialize_incoming_jsonrpc_message(line) {
|
||||
Ok(message) => Some(message),
|
||||
Err(error) => {
|
||||
debug!(
|
||||
"Failed to parse local MCP server message ({}): {error}",
|
||||
self.program_name
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Transport<RoleClient> for LocalStdioTransport {
|
||||
type Error = io::Error;
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "complete JSON-RPC frames must hold the stdin lock across writes"
|
||||
)]
|
||||
fn send(
|
||||
&mut self,
|
||||
item: TxJsonRpcMessage<RoleClient>,
|
||||
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
|
||||
let stdin = Arc::clone(&self.stdin);
|
||||
async move {
|
||||
let mut message = serde_json::to_vec(&item).map_err(io::Error::other)?;
|
||||
message.push(b'\n');
|
||||
let mut guard = stdin.lock().await;
|
||||
let stdin = guard
|
||||
.as_mut()
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::BrokenPipe, "MCP stdin closed"))?;
|
||||
stdin.write_all(&message).await?;
|
||||
stdin.flush().await
|
||||
}
|
||||
}
|
||||
|
||||
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<RoleClient>>> + Send {
|
||||
self.receive_message()
|
||||
}
|
||||
|
||||
async fn close(&mut self) -> Result<(), Self::Error> {
|
||||
self.stdin.lock().await.take();
|
||||
match tokio::time::timeout(Duration::from_secs(3), self.child.wait()).await {
|
||||
Ok(status) => {
|
||||
status?;
|
||||
Ok(())
|
||||
}
|
||||
Err(_) => self.child.kill().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
use std::ffi::OsStr;
|
||||
use std::io;
|
||||
|
||||
use rmcp::model::ProtocolVersion;
|
||||
use rmcp::service::ClientLifecycleMode;
|
||||
|
||||
/// MCP compatibility policy selected once when a Codex session is created.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum McpProtocolMode {
|
||||
/// Preserve the existing MCP initialization and OAuth behavior.
|
||||
#[default]
|
||||
Legacy,
|
||||
/// Allow the MCP 2026-07-28 discovery and request lifecycle.
|
||||
V20260728,
|
||||
}
|
||||
|
||||
impl McpProtocolMode {
|
||||
/// Returns the newest protocol version this compatibility policy can use.
|
||||
pub fn preferred_protocol_version(self) -> ProtocolVersion {
|
||||
match self {
|
||||
Self::Legacy => ProtocolVersion::V_2025_06_18,
|
||||
Self::V20260728 => ProtocolVersion::V_2026_07_28,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn client_lifecycle(self) -> ClientLifecycleMode {
|
||||
match self {
|
||||
Self::Legacy => ClientLifecycleMode::Initialize,
|
||||
Self::V20260728 => ClientLifecycleMode::Auto {
|
||||
preferred_versions: vec![ProtocolVersion::V_2026_07_28],
|
||||
legacy_version: Some(ProtocolVersion::V_2025_06_18),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn stdio_mode(self, requested_version: Option<&OsStr>) -> io::Result<Self> {
|
||||
match (self, requested_version) {
|
||||
(Self::Legacy, _) => Ok(Self::Legacy),
|
||||
(_, None) => Ok(Self::Legacy),
|
||||
(Self::V20260728, Some(version)) if version == OsStr::new("2026-07-28") => {
|
||||
Ok(Self::V20260728)
|
||||
}
|
||||
(_, Some(version)) => Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"unsupported CODEX_MCP_PROTOCOL_VERSION `{}` for stdio MCP server; expected `2026-07-28`",
|
||||
version.to_string_lossy()
|
||||
),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn protocol_modes_select_compatible_sdk_lifecycles() {
|
||||
assert_eq!(
|
||||
McpProtocolMode::Legacy.preferred_protocol_version(),
|
||||
ProtocolVersion::V_2025_06_18
|
||||
);
|
||||
assert_eq!(
|
||||
McpProtocolMode::Legacy.client_lifecycle(),
|
||||
ClientLifecycleMode::Initialize
|
||||
);
|
||||
assert_eq!(
|
||||
McpProtocolMode::V20260728.preferred_protocol_version(),
|
||||
ProtocolVersion::V_2026_07_28
|
||||
);
|
||||
assert_eq!(
|
||||
McpProtocolMode::V20260728.client_lifecycle(),
|
||||
ClientLifecycleMode::Auto {
|
||||
preferred_versions: vec![ProtocolVersion::V_2026_07_28],
|
||||
legacy_version: Some(ProtocolVersion::V_2025_06_18),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stdio_requires_both_the_modern_feature_and_a_server_opt_in() {
|
||||
let modern_version = Some(OsStr::new("2026-07-28"));
|
||||
|
||||
assert_eq!(
|
||||
McpProtocolMode::Legacy
|
||||
.stdio_mode(/*requested_version*/ None)
|
||||
.unwrap(),
|
||||
McpProtocolMode::Legacy
|
||||
);
|
||||
assert_eq!(
|
||||
McpProtocolMode::Legacy.stdio_mode(modern_version).unwrap(),
|
||||
McpProtocolMode::Legacy
|
||||
);
|
||||
assert_eq!(
|
||||
McpProtocolMode::V20260728
|
||||
.stdio_mode(/*requested_version*/ None)
|
||||
.unwrap(),
|
||||
McpProtocolMode::Legacy
|
||||
);
|
||||
assert_eq!(
|
||||
McpProtocolMode::V20260728
|
||||
.stdio_mode(modern_version)
|
||||
.unwrap(),
|
||||
McpProtocolMode::V20260728
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stdio_rejects_unknown_protocol_markers() {
|
||||
let error = McpProtocolMode::V20260728
|
||||
.stdio_mode(Some(OsStr::new("1999-01-01")))
|
||||
.expect_err("an unknown protocol marker must fail");
|
||||
|
||||
assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
|
||||
assert!(error.to_string().contains("1999-01-01"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_stdio_does_not_interpret_existing_protocol_markers() {
|
||||
assert_eq!(
|
||||
McpProtocolMode::Legacy
|
||||
.stdio_mode(Some(OsStr::new("1999-01-01")))
|
||||
.unwrap(),
|
||||
McpProtocolMode::Legacy
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsStr;
|
||||
use std::ffi::OsString;
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
@@ -45,9 +46,9 @@ use rmcp::model::RequestParamsMeta;
|
||||
use rmcp::model::ServerResult;
|
||||
use rmcp::model::Tool;
|
||||
use rmcp::service::ClientCacheConfig;
|
||||
use rmcp::service::ClientServiceExt;
|
||||
use rmcp::service::RoleClient;
|
||||
use rmcp::service::RunningService;
|
||||
use rmcp::service::{self};
|
||||
use rmcp::transport::AuthorizationManager;
|
||||
use rmcp::transport::StreamableHttpClientTransport;
|
||||
use rmcp::transport::auth::AuthClient;
|
||||
@@ -75,6 +76,7 @@ use crate::oauth::ResolvedOAuthTokens;
|
||||
use crate::oauth::StoredOAuthTokens;
|
||||
use crate::oauth::resolve_oauth_tokens_from_store_policy;
|
||||
use crate::oauth_http_client::OAuthHttpClientAdapter;
|
||||
use crate::protocol_mode::McpProtocolMode;
|
||||
use crate::stdio_server_launcher::StdioServerCommand;
|
||||
use crate::stdio_server_launcher::StdioServerLauncher;
|
||||
use crate::stdio_server_launcher::StdioServerProcessHandle;
|
||||
@@ -328,12 +330,18 @@ pub struct RmcpClient {
|
||||
state: Mutex<ClientState>,
|
||||
stdio_process: Option<StdioServerProcessHandle>,
|
||||
transport_recipe: TransportRecipe,
|
||||
protocol_mode: McpProtocolMode,
|
||||
initialize_context: Mutex<Option<InitializeContext>>,
|
||||
session_recovery_lock: Semaphore,
|
||||
elicitation_pause_state: ElicitationPauseState,
|
||||
}
|
||||
|
||||
impl RmcpClient {
|
||||
/// Returns the protocol compatibility policy captured when this client was created.
|
||||
pub fn protocol_mode(&self) -> McpProtocolMode {
|
||||
self.protocol_mode
|
||||
}
|
||||
|
||||
pub async fn new_in_process_client(
|
||||
factory: Arc<dyn InProcessTransportFactory>,
|
||||
) -> io::Result<Self> {
|
||||
@@ -348,6 +356,7 @@ impl RmcpClient {
|
||||
}),
|
||||
stdio_process: None,
|
||||
transport_recipe,
|
||||
protocol_mode: McpProtocolMode::Legacy,
|
||||
initialize_context: Mutex::new(None),
|
||||
session_recovery_lock: Semaphore::new(/*permits*/ 1),
|
||||
elicitation_pause_state: ElicitationPauseState::new(),
|
||||
@@ -362,8 +371,45 @@ impl RmcpClient {
|
||||
cwd: Option<String>,
|
||||
launcher: Arc<dyn StdioServerLauncher>,
|
||||
) -> io::Result<Self> {
|
||||
Self::new_stdio_client_with_protocol_mode(
|
||||
program,
|
||||
args,
|
||||
env,
|
||||
env_vars,
|
||||
cwd,
|
||||
launcher,
|
||||
McpProtocolMode::Legacy,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Constructs a stdio client with an explicitly selected compatibility policy.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn new_stdio_client_with_protocol_mode(
|
||||
program: OsString,
|
||||
args: Vec<OsString>,
|
||||
mut env: Option<HashMap<OsString, OsString>>,
|
||||
env_vars: &[McpServerEnvVar],
|
||||
cwd: Option<String>,
|
||||
launcher: Arc<dyn StdioServerLauncher>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
) -> io::Result<Self> {
|
||||
let requested_stdio_version = match protocol_mode {
|
||||
McpProtocolMode::Legacy => None,
|
||||
McpProtocolMode::V20260728 => env
|
||||
.as_mut()
|
||||
.and_then(|env| env.remove(OsStr::new("CODEX_MCP_PROTOCOL_VERSION"))),
|
||||
};
|
||||
let protocol_mode = protocol_mode.stdio_mode(requested_stdio_version.as_deref())?;
|
||||
let transport_recipe = TransportRecipe::Stdio {
|
||||
command: StdioServerCommand::new(program, args, env, env_vars.to_vec(), cwd),
|
||||
command: StdioServerCommand::new(
|
||||
program,
|
||||
args,
|
||||
env,
|
||||
env_vars.to_vec(),
|
||||
cwd,
|
||||
protocol_mode,
|
||||
),
|
||||
launcher,
|
||||
};
|
||||
let transport = Self::create_pending_transport(&transport_recipe)
|
||||
@@ -382,6 +428,7 @@ impl RmcpClient {
|
||||
}),
|
||||
stdio_process,
|
||||
transport_recipe,
|
||||
protocol_mode,
|
||||
initialize_context: Mutex::new(None),
|
||||
session_recovery_lock: Semaphore::new(/*permits*/ 1),
|
||||
elicitation_pause_state: ElicitationPauseState::new(),
|
||||
@@ -399,6 +446,35 @@ impl RmcpClient {
|
||||
keyring_backend_kind: AuthKeyringBackendKind,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
auth_provider: Option<SharedAuthProvider>,
|
||||
) -> Result<Self> {
|
||||
Self::new_streamable_http_client_with_protocol_mode(
|
||||
server_name,
|
||||
url,
|
||||
bearer_token,
|
||||
http_headers,
|
||||
env_http_headers,
|
||||
store_mode,
|
||||
keyring_backend_kind,
|
||||
http_client,
|
||||
auth_provider,
|
||||
McpProtocolMode::Legacy,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Constructs a streamable HTTP client with an explicitly selected compatibility policy.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn new_streamable_http_client_with_protocol_mode(
|
||||
server_name: &str,
|
||||
url: &str,
|
||||
bearer_token: Option<String>,
|
||||
http_headers: Option<HashMap<String, String>>,
|
||||
env_http_headers: Option<HashMap<String, String>>,
|
||||
store_mode: OAuthCredentialsStoreMode,
|
||||
keyring_backend_kind: AuthKeyringBackendKind,
|
||||
http_client: Arc<dyn HttpClient>,
|
||||
auth_provider: Option<SharedAuthProvider>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
) -> Result<Self> {
|
||||
let transport_recipe = TransportRecipe::StreamableHttp {
|
||||
server_name: server_name.to_string(),
|
||||
@@ -419,6 +495,7 @@ impl RmcpClient {
|
||||
}),
|
||||
stdio_process: None,
|
||||
transport_recipe,
|
||||
protocol_mode,
|
||||
initialize_context: Mutex::new(None),
|
||||
session_recovery_lock: Semaphore::new(/*permits*/ 1),
|
||||
elicitation_pause_state: ElicitationPauseState::new(),
|
||||
@@ -933,6 +1010,7 @@ impl RmcpClient {
|
||||
}
|
||||
|
||||
async fn connect_pending_transport(
|
||||
&self,
|
||||
pending_transport: PendingTransport,
|
||||
client_service: ElicitationClientService,
|
||||
timeout: Option<Duration>,
|
||||
@@ -940,24 +1018,33 @@ impl RmcpClient {
|
||||
Arc<RunningService<RoleClient, ElicitationClientService>>,
|
||||
Option<OAuthPersistor>,
|
||||
)> {
|
||||
let lifecycle = self.protocol_mode.client_lifecycle();
|
||||
let (transport, oauth_persistor) = match pending_transport {
|
||||
PendingTransport::InProcess { transport } => (
|
||||
service::serve_client(client_service, transport).boxed(),
|
||||
client_service
|
||||
.serve_with_lifecycle(transport, lifecycle)
|
||||
.boxed(),
|
||||
None,
|
||||
),
|
||||
PendingTransport::Stdio { transport } => (
|
||||
service::serve_client(client_service, *transport).boxed(),
|
||||
client_service
|
||||
.serve_with_lifecycle(*transport, lifecycle)
|
||||
.boxed(),
|
||||
None,
|
||||
),
|
||||
PendingTransport::StreamableHttp { transport } => (
|
||||
service::serve_client(client_service, transport).boxed(),
|
||||
client_service
|
||||
.serve_with_lifecycle(transport, lifecycle)
|
||||
.boxed(),
|
||||
None,
|
||||
),
|
||||
PendingTransport::StreamableHttpWithOAuth {
|
||||
transport,
|
||||
oauth_persistor,
|
||||
} => (
|
||||
service::serve_client(client_service, transport).boxed(),
|
||||
client_service
|
||||
.serve_with_lifecycle(transport, lifecycle)
|
||||
.boxed(),
|
||||
Some(oauth_persistor),
|
||||
),
|
||||
};
|
||||
|
||||
@@ -56,7 +56,9 @@ use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::executor_process_transport::ExecutorProcessTransport;
|
||||
use crate::local_stdio_transport::LocalStdioTransport;
|
||||
use crate::program_resolver;
|
||||
use crate::protocol_mode::McpProtocolMode;
|
||||
use crate::utils::create_env_for_mcp_server;
|
||||
use crate::utils::create_env_overlay_for_remote_mcp_server;
|
||||
use crate::utils::remote_mcp_env_var_names;
|
||||
@@ -85,6 +87,7 @@ pub struct StdioServerCommand {
|
||||
env: Option<HashMap<OsString, OsString>>,
|
||||
env_vars: Vec<McpServerEnvVar>,
|
||||
cwd: Option<String>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
}
|
||||
|
||||
/// Client-side rmcp transport for a launched MCP stdio server.
|
||||
@@ -98,7 +101,8 @@ pub struct StdioServerTransport {
|
||||
}
|
||||
|
||||
enum StdioServerTransportInner {
|
||||
Local(TokioChildProcess),
|
||||
LocalLegacy(TokioChildProcess),
|
||||
LocalModern(LocalStdioTransport),
|
||||
Executor(ExecutorProcessTransport),
|
||||
}
|
||||
|
||||
@@ -113,7 +117,8 @@ impl Transport<RoleClient> for StdioServerTransport {
|
||||
// wrapper keeps process placement private while leaving rmcp's send
|
||||
// semantics unchanged.
|
||||
match &mut self.inner {
|
||||
StdioServerTransportInner::Local(transport) => transport.send(item).boxed(),
|
||||
StdioServerTransportInner::LocalLegacy(transport) => transport.send(item).boxed(),
|
||||
StdioServerTransportInner::LocalModern(transport) => transport.send(item).boxed(),
|
||||
StdioServerTransportInner::Executor(transport) => transport.send(item).boxed(),
|
||||
}
|
||||
}
|
||||
@@ -123,7 +128,8 @@ impl Transport<RoleClient> for StdioServerTransport {
|
||||
// executor variant turns pushed process-output events back into the
|
||||
// line-delimited JSON stream expected by rmcp.
|
||||
match &mut self.inner {
|
||||
StdioServerTransportInner::Local(transport) => transport.receive().boxed(),
|
||||
StdioServerTransportInner::LocalLegacy(transport) => transport.receive().boxed(),
|
||||
StdioServerTransportInner::LocalModern(transport) => transport.receive().boxed(),
|
||||
StdioServerTransportInner::Executor(transport) => transport.receive().boxed(),
|
||||
}
|
||||
}
|
||||
@@ -131,7 +137,8 @@ impl Transport<RoleClient> for StdioServerTransport {
|
||||
async fn close(&mut self) -> std::result::Result<(), Self::Error> {
|
||||
self.process.terminate().await?;
|
||||
match &mut self.inner {
|
||||
StdioServerTransportInner::Local(transport) => transport.close().await,
|
||||
StdioServerTransportInner::LocalLegacy(transport) => transport.close().await,
|
||||
StdioServerTransportInner::LocalModern(transport) => transport.close().await,
|
||||
StdioServerTransportInner::Executor(transport) => transport.close().await,
|
||||
}
|
||||
}
|
||||
@@ -152,6 +159,7 @@ impl StdioServerCommand {
|
||||
env: Option<HashMap<OsString, OsString>>,
|
||||
env_vars: Vec<McpServerEnvVar>,
|
||||
cwd: Option<String>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
) -> Self {
|
||||
Self {
|
||||
program,
|
||||
@@ -159,6 +167,7 @@ impl StdioServerCommand {
|
||||
env,
|
||||
env_vars,
|
||||
cwd,
|
||||
protocol_mode,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -253,6 +262,7 @@ impl LocalStdioServerLauncher {
|
||||
env,
|
||||
env_vars,
|
||||
cwd,
|
||||
protocol_mode,
|
||||
} = command;
|
||||
let program_name = program.to_string_lossy().into_owned();
|
||||
let envs = create_env_for_mcp_server(env, &env_vars).map_err(io::Error::other)?;
|
||||
@@ -272,12 +282,32 @@ impl LocalStdioServerLauncher {
|
||||
#[cfg(unix)]
|
||||
command.process_group(0);
|
||||
|
||||
let (transport, stderr) = TokioChildProcess::builder(command)
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()?;
|
||||
let (transport, stderr, process_id) = match protocol_mode {
|
||||
McpProtocolMode::Legacy => {
|
||||
let (transport, stderr) = TokioChildProcess::builder(command)
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()?;
|
||||
let process_id = transport.id();
|
||||
(
|
||||
StdioServerTransportInner::LocalLegacy(transport),
|
||||
stderr,
|
||||
process_id,
|
||||
)
|
||||
}
|
||||
McpProtocolMode::V20260728 => {
|
||||
let (transport, stderr) =
|
||||
LocalStdioTransport::spawn(command, program_name.clone())?;
|
||||
let process_id = transport.id();
|
||||
(
|
||||
StdioServerTransportInner::LocalModern(transport),
|
||||
stderr,
|
||||
process_id,
|
||||
)
|
||||
}
|
||||
};
|
||||
let process = StdioServerProcessHandle::local(
|
||||
program_name.clone(),
|
||||
transport.id().map(LocalProcessTerminator::new),
|
||||
process_id.map(LocalProcessTerminator::new),
|
||||
);
|
||||
|
||||
if let Some(stderr) = stderr {
|
||||
@@ -299,7 +329,7 @@ impl LocalStdioServerLauncher {
|
||||
}
|
||||
|
||||
Ok(StdioServerTransport {
|
||||
inner: StdioServerTransportInner::Local(transport),
|
||||
inner: transport,
|
||||
process,
|
||||
})
|
||||
}
|
||||
@@ -482,6 +512,7 @@ impl ExecutorStdioServerLauncher {
|
||||
env,
|
||||
env_vars,
|
||||
cwd,
|
||||
protocol_mode: _,
|
||||
} = command;
|
||||
let Some(cwd) = cwd else {
|
||||
return Err(io::Error::other(
|
||||
|
||||
@@ -76,12 +76,9 @@ impl RmcpClient {
|
||||
}
|
||||
let attempt_timeout = remaining_initialize_timeout(timeout, retry_deadline)?;
|
||||
|
||||
match Self::connect_pending_transport(
|
||||
transport,
|
||||
client_service.clone(),
|
||||
attempt_timeout,
|
||||
)
|
||||
.await
|
||||
match self
|
||||
.connect_pending_transport(transport, client_service.clone(), attempt_timeout)
|
||||
.await
|
||||
{
|
||||
Ok(result) => return Ok(result),
|
||||
Err(error) if should_retry && Self::is_retryable_initialize_error(&error) => {
|
||||
@@ -124,7 +121,10 @@ impl RmcpClient {
|
||||
fn is_retryable_client_initialize_error(error: &rmcp::service::ClientInitializeError) -> bool {
|
||||
match error {
|
||||
rmcp::service::ClientInitializeError::TransportError { error, context }
|
||||
if context.as_ref() == "send initialize request" =>
|
||||
if matches!(
|
||||
context.as_ref(),
|
||||
"send initialize request" | "send discover request"
|
||||
) =>
|
||||
{
|
||||
error
|
||||
.error
|
||||
|
||||
@@ -10,8 +10,9 @@ use crate::http_client_adapter::StreamableHttpClientAdapterError;
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn retryable_initialize_error_includes_initialized_notification_context() {
|
||||
fn retryable_initialize_error_includes_discovery_and_initialized_notification_context() {
|
||||
let contexts = [
|
||||
"send discover request",
|
||||
"send initialize request",
|
||||
"send initialized notification",
|
||||
"receive initialize response",
|
||||
@@ -21,7 +22,7 @@ fn retryable_initialize_error_includes_initialized_notification_context() {
|
||||
contexts.map(|context| {
|
||||
RmcpClient::is_retryable_client_initialize_error(&retryable_initialize_error(context))
|
||||
}),
|
||||
[true, true, false],
|
||||
[true, true, true, false],
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,203 @@
|
||||
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::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::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 MODERN_VERSION: &str = "2026-07-28";
|
||||
const LEGACY_VERSION: &str = "2025-06-18";
|
||||
|
||||
async fn initialize_modern_client(server: &MockServer) -> anyhow::Result<RmcpClient> {
|
||||
let client = RmcpClient::new_streamable_http_client_with_protocol_mode(
|
||||
"sse-discovery-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 params = InitializeRequestParams::new(
|
||||
ClientCapabilities::default(),
|
||||
Implementation::new("codex-sse-discovery-test", "0.0.0"),
|
||||
)
|
||||
.with_protocol_version(ProtocolVersion::V_2025_06_18);
|
||||
|
||||
client
|
||||
.initialize(
|
||||
params,
|
||||
Some(Duration::from_secs(5)),
|
||||
Box::new(|_, _| {
|
||||
async {
|
||||
Ok(ElicitationResponse {
|
||||
action: ElicitationAction::Accept,
|
||||
content: Some(json!({})),
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
fn sse_response(message: Value) -> ResponseTemplate {
|
||||
ResponseTemplate::new(200).set_body_raw(
|
||||
format!("event: message\ndata: {message}\n\n"),
|
||||
"text/event-stream",
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn modern_sse_discovery_accepts_metadata_namespaced_server_identity() -> 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");
|
||||
|
||||
sse_response(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"resultType": "complete",
|
||||
"supportedVersions": [MODERN_VERSION],
|
||||
"capabilities": {"tools": {}},
|
||||
"_meta": {
|
||||
"io.modelcontextprotocol/serverInfo": {
|
||||
"name": "modern-sse-test",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
},
|
||||
"ttlMs": 0,
|
||||
"cacheScope": "private",
|
||||
},
|
||||
}))
|
||||
})
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = initialize_modern_client(&server).await?;
|
||||
client.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn modern_sse_discovery_falls_back_for_correlated_method_not_found() -> anyhow::Result<()> {
|
||||
let server = MockServer::start().await;
|
||||
let observed = Arc::new(Mutex::new(Vec::<String>::new()));
|
||||
let recorded = Arc::clone(&observed);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/mcp"))
|
||||
.respond_with(move |request: &Request| {
|
||||
let body: Value = request.body_json().expect("valid JSON-RPC request");
|
||||
let method = body["method"].as_str().expect("JSON-RPC method");
|
||||
recorded.lock().expect("requests lock").push(method.into());
|
||||
|
||||
match method {
|
||||
"server/discover" => sse_response(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"error": {"code": -32601, "message": "method not found"},
|
||||
})),
|
||||
"initialize" => {
|
||||
assert_eq!(body["params"]["protocolVersion"], LEGACY_VERSION);
|
||||
ResponseTemplate::new(200).set_body_json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"],
|
||||
"result": {
|
||||
"protocolVersion": LEGACY_VERSION,
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {
|
||||
"name": "legacy-sse-test",
|
||||
"version": "1.0.0",
|
||||
},
|
||||
},
|
||||
}))
|
||||
}
|
||||
"notifications/initialized" => ResponseTemplate::new(202),
|
||||
other => panic!("unexpected legacy SSE fallback request: {other}"),
|
||||
}
|
||||
})
|
||||
.expect(3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let client = initialize_modern_client(&server).await?;
|
||||
assert_eq!(
|
||||
*observed.lock().expect("requests lock"),
|
||||
vec!["server/discover", "initialize", "notifications/initialized"]
|
||||
);
|
||||
client.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn modern_sse_discovery_rejects_uncorrelated_method_not_found() -> anyhow::Result<()> {
|
||||
for (case, rejected_id) in [("null", json!(null)), ("mismatched", json!("unrelated"))] {
|
||||
let server = MockServer::start().await;
|
||||
let observed = Arc::new(Mutex::new(Vec::<String>::new()));
|
||||
let recorded = Arc::clone(&observed);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/mcp"))
|
||||
.respond_with(move |request: &Request| {
|
||||
let body: Value = request.body_json().expect("valid JSON-RPC request");
|
||||
let method = body["method"].as_str().expect("JSON-RPC method");
|
||||
recorded.lock().expect("requests lock").push(method.into());
|
||||
assert_eq!(method, "server/discover");
|
||||
|
||||
sse_response(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": rejected_id.clone(),
|
||||
"error": {"code": -32601, "message": "method not found"},
|
||||
}))
|
||||
})
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
initialize_modern_client(&server).await.is_err(),
|
||||
"{case} SSE discovery error must not downgrade to legacy"
|
||||
);
|
||||
assert_eq!(
|
||||
*observed.lock().expect("requests lock"),
|
||||
vec!["server/discover"],
|
||||
"{case} SSE discovery error must not initialize a legacy session"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsString;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
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::json;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn modern_local_and_executor_stdio_discover_metadata_identity_and_catalogs()
|
||||
-> anyhow::Result<()> {
|
||||
let server = codex_utils_cargo_bin::cargo_bin("test_mcp_2026_discovery_stdio_server")?;
|
||||
|
||||
for executor in [false, true] {
|
||||
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 client = RmcpClient::new_stdio_client_with_protocol_mode(
|
||||
server.clone().into(),
|
||||
Vec::new(),
|
||||
Some(HashMap::from([(
|
||||
OsString::from("CODEX_MCP_PROTOCOL_VERSION"),
|
||||
OsString::from("2026-07-28"),
|
||||
)])),
|
||||
&[],
|
||||
Some(std::env::current_dir()?.to_string_lossy().into_owned()),
|
||||
launcher,
|
||||
McpProtocolMode::V20260728,
|
||||
)
|
||||
.await?;
|
||||
client
|
||||
.initialize(
|
||||
InitializeRequestParams::new(
|
||||
ClientCapabilities::default(),
|
||||
Implementation::new("stdio-discovery-test", "1.0.0"),
|
||||
)
|
||||
.with_protocol_version(ProtocolVersion::V_2025_06_18),
|
||||
Some(Duration::from_secs(10)),
|
||||
Box::new(|_, _| {
|
||||
async {
|
||||
Ok(ElicitationResponse {
|
||||
action: ElicitationAction::Decline,
|
||||
content: None,
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let tools = client
|
||||
.list_tools(/*params*/ None, Some(Duration::from_secs(10)))
|
||||
.await?;
|
||||
assert_eq!(tools.tools[0].name.as_ref(), "stdio_echo");
|
||||
let resources = client
|
||||
.list_resources(/*params*/ None, Some(Duration::from_secs(10)))
|
||||
.await?;
|
||||
assert_eq!(resources.resources[0].uri, "test://stdio/resource");
|
||||
let templates = client
|
||||
.list_resource_templates(/*params*/ None, Some(Duration::from_secs(10)))
|
||||
.await?;
|
||||
assert_eq!(
|
||||
templates.resource_templates[0].name,
|
||||
"stdio resource template"
|
||||
);
|
||||
client.shutdown().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn legacy_stdio_preserves_existing_protocol_marker_environment() -> anyhow::Result<()> {
|
||||
let server = codex_utils_cargo_bin::cargo_bin("test_stdio_server")?;
|
||||
|
||||
for executor in [false, true] {
|
||||
for version in ["2026-07-28", "1999-01-01"] {
|
||||
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 client = RmcpClient::new_stdio_client_with_protocol_mode(
|
||||
server.clone().into(),
|
||||
Vec::new(),
|
||||
Some(HashMap::from([(
|
||||
OsString::from("CODEX_MCP_PROTOCOL_VERSION"),
|
||||
OsString::from(version),
|
||||
)])),
|
||||
&[],
|
||||
Some(std::env::current_dir()?.to_string_lossy().into_owned()),
|
||||
launcher,
|
||||
McpProtocolMode::Legacy,
|
||||
)
|
||||
.await?;
|
||||
|
||||
client
|
||||
.initialize(
|
||||
InitializeRequestParams::new(
|
||||
ClientCapabilities::default(),
|
||||
Implementation::new("stdio-legacy-environment-test", "1.0.0"),
|
||||
)
|
||||
.with_protocol_version(ProtocolVersion::V_2025_06_18),
|
||||
Some(Duration::from_secs(10)),
|
||||
Box::new(|_, _| {
|
||||
async {
|
||||
Ok(ElicitationResponse {
|
||||
action: ElicitationAction::Decline,
|
||||
content: None,
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let result = client
|
||||
.call_tool(
|
||||
"echo".to_string(),
|
||||
Some(json!({
|
||||
"message": "legacy environment",
|
||||
"env_var": "CODEX_MCP_PROTOCOL_VERSION",
|
||||
})),
|
||||
/*meta*/ None,
|
||||
Some(Duration::from_secs(10)),
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(
|
||||
result.structured_content,
|
||||
Some(json!({
|
||||
"echo": "ECHOING: legacy environment",
|
||||
"env": version,
|
||||
}))
|
||||
);
|
||||
client.shutdown().await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsString;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
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;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stdio_message_limits_preserve_legacy_local_compatibility() -> anyhow::Result<()> {
|
||||
let server = codex_utils_cargo_bin::cargo_bin("test_stdio_server")?;
|
||||
|
||||
for (executor, protocol_mode, accepts_oversized) in [
|
||||
(false, McpProtocolMode::Legacy, true),
|
||||
(false, McpProtocolMode::V20260728, false),
|
||||
(true, McpProtocolMode::Legacy, false),
|
||||
] {
|
||||
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 mut env = HashMap::from([(
|
||||
OsString::from("MCP_TEST_OVERSIZED_TOOL_DESCRIPTION"),
|
||||
OsString::from("1"),
|
||||
)]);
|
||||
if protocol_mode == McpProtocolMode::V20260728 {
|
||||
env.insert(
|
||||
OsString::from("CODEX_MCP_PROTOCOL_VERSION"),
|
||||
OsString::from("2026-07-28"),
|
||||
);
|
||||
}
|
||||
let client = RmcpClient::new_stdio_client_with_protocol_mode(
|
||||
server.clone().into(),
|
||||
Vec::new(),
|
||||
Some(env),
|
||||
&[],
|
||||
Some(std::env::current_dir()?.to_string_lossy().into_owned()),
|
||||
launcher,
|
||||
protocol_mode,
|
||||
)
|
||||
.await?;
|
||||
client
|
||||
.initialize(
|
||||
InitializeRequestParams::new(
|
||||
ClientCapabilities::default(),
|
||||
Implementation::new("stdio-limit-test", "1.0.0"),
|
||||
)
|
||||
.with_protocol_version(ProtocolVersion::V_2025_06_18),
|
||||
Some(Duration::from_secs(10)),
|
||||
Box::new(|_, _| {
|
||||
async {
|
||||
Ok(ElicitationResponse {
|
||||
action: ElicitationAction::Decline,
|
||||
content: None,
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
.boxed()
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let result = client
|
||||
.list_tools(/*params*/ None, Some(Duration::from_secs(10)))
|
||||
.await;
|
||||
assert_eq!(
|
||||
result.is_ok(),
|
||||
accepts_oversized,
|
||||
"unexpected stdio size handling (executor={executor}, mode={protocol_mode:?})"
|
||||
);
|
||||
client.shutdown().await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user