mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Add supports_parallel_tool_calls flag to included mcps (#17667)
## Why
For more advanced MCP usage, we want the model to be able to emit
parallel MCP tool calls and have Codex execute eligible ones
concurrently, instead of forcing all MCP calls through the serial block.
The main design choice was where to thread the config. I made this
server-level because parallel safety depends on the MCP server
implementation. Codex reads the flag from `mcp_servers`, threads the
opted-in server names into `ToolRouter`, and checks the parsed
`ToolPayload::Mcp { server, .. }` at execution time. That avoids relying
on model-visible tool names, which can be incomplete in
deferred/search-tool paths or ambiguous for similarly named
servers/tools.
## What was added
Added `supports_parallel_tool_calls` for MCP servers.
Before:
```toml
[mcp_servers.docs]
command = "docs-server"
```
After:
```toml
[mcp_servers.docs]
command = "docs-server"
supports_parallel_tool_calls = true
```
MCP calls remain serial by default. Only tools from opted-in servers are
eligible to run in parallel. Docs also now warn to enable this only when
the server’s tools are safe to run concurrently, especially around
shared state or read/write races.
## Testing
Tested with a local stdio MCP server exposing real delay tools. The
model/Responses side was mocked only to deterministically emit two MCP
calls in the same turn.
Each test called `query_with_delay` and `query_with_delay_2` with `{
"seconds": 25 }`.
| Build/config | Observed | Wall time |
| --- | --- | --- |
| main with flag enabled | serial | `58.79s` |
| PR with flag enabled | parallel | `31.73s` |
| PR without flag | serial | `56.70s` |
PR with flag enabled showed both tools start before either completed;
main and PR-without-flag completed the first delay before starting the
second.
Also added an integration test.
Additional checks:
- `cargo test -p codex-tools` passed
- `cargo test -p codex-core
mcp_parallel_support_uses_exact_payload_server` passed
- `git diff --check` passed
This commit is contained in:
@@ -1,6 +1,9 @@
|
||||
use std::borrow::Cow;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::hash_map::Entry;
|
||||
use std::sync::Arc;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use rmcp::ErrorData as McpError;
|
||||
use rmcp::ServiceExt;
|
||||
@@ -25,7 +28,9 @@ use rmcp::model::Tool;
|
||||
use rmcp::model::ToolAnnotations;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
use tokio::sync::Barrier;
|
||||
use tokio::task;
|
||||
use tokio::time::sleep;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestToolServer {
|
||||
@@ -47,6 +52,7 @@ impl TestToolServer {
|
||||
let tools = vec![
|
||||
Self::echo_tool(),
|
||||
Self::echo_dash_tool(),
|
||||
Self::sync_tool(),
|
||||
Self::image_tool(),
|
||||
Self::image_scenario_tool(),
|
||||
];
|
||||
@@ -112,6 +118,50 @@ impl TestToolServer {
|
||||
tool
|
||||
}
|
||||
|
||||
fn sync_tool() -> Tool {
|
||||
#[expect(clippy::expect_used)]
|
||||
let schema: JsonObject = serde_json::from_value(json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"sleep_before_ms": { "type": "number" },
|
||||
"sleep_after_ms": { "type": "number" },
|
||||
"barrier": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": { "type": "string" },
|
||||
"participants": { "type": "number" },
|
||||
"timeout_ms": { "type": "number" }
|
||||
},
|
||||
"required": ["id", "participants"],
|
||||
"additionalProperties": false
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}))
|
||||
.expect("sync tool schema should deserialize");
|
||||
|
||||
let mut tool = Tool::new(
|
||||
Cow::Borrowed("sync"),
|
||||
Cow::Borrowed(
|
||||
"Synchronize concurrent test calls and optionally delay before or after the barrier.",
|
||||
),
|
||||
Arc::new(schema),
|
||||
);
|
||||
#[expect(clippy::expect_used)]
|
||||
let output_schema: JsonObject = serde_json::from_value(json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"result": { "type": "string" }
|
||||
},
|
||||
"required": ["result"],
|
||||
"additionalProperties": false
|
||||
}))
|
||||
.expect("sync tool output schema should deserialize");
|
||||
tool.output_schema = Some(Arc::new(output_schema));
|
||||
tool.annotations = Some(ToolAnnotations::new().read_only(true));
|
||||
tool
|
||||
}
|
||||
|
||||
fn image_tool() -> Tool {
|
||||
#[expect(clippy::expect_used)]
|
||||
let schema: JsonObject = serde_json::from_value(serde_json::json!({
|
||||
@@ -227,6 +277,42 @@ struct EchoArgs {
|
||||
env_var: Option<String>,
|
||||
}
|
||||
|
||||
const DEFAULT_SYNC_TIMEOUT_MS: u64 = 1_000;
|
||||
|
||||
static SYNC_BARRIERS: OnceLock<tokio::sync::Mutex<HashMap<String, SyncBarrierState>>> =
|
||||
OnceLock::new();
|
||||
|
||||
struct SyncBarrierState {
|
||||
barrier: Arc<Barrier>,
|
||||
participants: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SyncBarrierArgs {
|
||||
id: String,
|
||||
participants: usize,
|
||||
#[serde(default = "default_sync_timeout_ms")]
|
||||
timeout_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SyncArgs {
|
||||
#[serde(default)]
|
||||
sleep_before_ms: Option<u64>,
|
||||
#[serde(default)]
|
||||
sleep_after_ms: Option<u64>,
|
||||
#[serde(default)]
|
||||
barrier: Option<SyncBarrierArgs>,
|
||||
}
|
||||
|
||||
fn default_sync_timeout_ms() -> u64 {
|
||||
DEFAULT_SYNC_TIMEOUT_MS
|
||||
}
|
||||
|
||||
fn sync_barrier_map() -> &'static tokio::sync::Mutex<HashMap<String, SyncBarrierState>> {
|
||||
SYNC_BARRIERS.get_or_init(|| tokio::sync::Mutex::new(HashMap::new()))
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
/// Scenarios for `image_scenario`, intended to exercise Codex TUI handling of MCP image outputs.
|
||||
@@ -387,6 +473,10 @@ impl ServerHandler for TestToolServer {
|
||||
let args = Self::parse_call_args::<ImageScenarioArgs>(&request, "image_scenario")?;
|
||||
Self::image_scenario_result(args)
|
||||
}
|
||||
"sync" => {
|
||||
let args = Self::parse_call_args::<SyncArgs>(&request, "sync")?;
|
||||
Self::sync_result(args).await
|
||||
}
|
||||
other => Err(McpError::invalid_params(
|
||||
format!("unknown tool: {other}"),
|
||||
None,
|
||||
@@ -469,6 +559,102 @@ impl TestToolServer {
|
||||
|
||||
Ok(CallToolResult::success(content))
|
||||
}
|
||||
|
||||
async fn sync_result(args: SyncArgs) -> Result<CallToolResult, McpError> {
|
||||
if let Some(delay) = args.sleep_before_ms
|
||||
&& delay > 0
|
||||
{
|
||||
sleep(Duration::from_millis(delay)).await;
|
||||
}
|
||||
|
||||
if let Some(barrier) = args.barrier {
|
||||
wait_on_sync_barrier(barrier).await?;
|
||||
}
|
||||
|
||||
if let Some(delay) = args.sleep_after_ms
|
||||
&& delay > 0
|
||||
{
|
||||
sleep(Duration::from_millis(delay)).await;
|
||||
}
|
||||
|
||||
Ok(CallToolResult {
|
||||
content: Vec::new(),
|
||||
structured_content: Some(json!({ "result": "ok" })),
|
||||
is_error: Some(false),
|
||||
meta: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_on_sync_barrier(args: SyncBarrierArgs) -> Result<(), McpError> {
|
||||
if args.participants == 0 {
|
||||
return Err(McpError::invalid_params(
|
||||
"barrier participants must be greater than zero",
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
if args.timeout_ms == 0 {
|
||||
return Err(McpError::invalid_params(
|
||||
"barrier timeout must be greater than zero",
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
let barrier_id = args.id.clone();
|
||||
let barrier = {
|
||||
let mut map = sync_barrier_map().lock().await;
|
||||
match map.entry(barrier_id.clone()) {
|
||||
Entry::Occupied(entry) => {
|
||||
let state = entry.get();
|
||||
if state.participants != args.participants {
|
||||
let existing = state.participants;
|
||||
return Err(McpError::invalid_params(
|
||||
format!(
|
||||
"barrier {barrier_id} already registered with {existing} participants"
|
||||
),
|
||||
None,
|
||||
));
|
||||
}
|
||||
state.barrier.clone()
|
||||
}
|
||||
Entry::Vacant(entry) => {
|
||||
let barrier = Arc::new(Barrier::new(args.participants));
|
||||
entry.insert(SyncBarrierState {
|
||||
barrier: barrier.clone(),
|
||||
participants: args.participants,
|
||||
});
|
||||
barrier
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let wait_result =
|
||||
match tokio::time::timeout(Duration::from_millis(args.timeout_ms), barrier.wait()).await {
|
||||
Ok(wait_result) => wait_result,
|
||||
Err(_) => {
|
||||
remove_sync_barrier_if_current(&barrier_id, &barrier).await;
|
||||
return Err(McpError::invalid_params(
|
||||
"sync barrier wait timed out",
|
||||
None,
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
if wait_result.is_leader() {
|
||||
remove_sync_barrier_if_current(&barrier_id, &barrier).await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn remove_sync_barrier_if_current(barrier_id: &str, barrier: &Arc<Barrier>) {
|
||||
let mut map = sync_barrier_map().lock().await;
|
||||
if let Some(state) = map.get(barrier_id)
|
||||
&& Arc::ptr_eq(&state.barrier, barrier)
|
||||
{
|
||||
map.remove(barrier_id);
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_data_url(url: &str) -> Option<(String, String)> {
|
||||
|
||||
Reference in New Issue
Block a user