diff --git a/codex-rs/mcp-server/tests/codex_tool.rs b/codex-rs/mcp-server/tests/codex_tool.rs index 92f11eaa..1ebd10a7 100644 --- a/codex-rs/mcp-server/tests/codex_tool.rs +++ b/codex-rs/mcp-server/tests/codex_tool.rs @@ -35,7 +35,7 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs /// Test that a shell command that is not on the "trusted" list triggers an /// elicitation request to the MCP and that sending the approval runs the /// command, as expected. -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_shell_command_approval_triggers_elicitation() { if env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { println!( @@ -114,6 +114,16 @@ async fn shell_command_approval_triggers_elicitation() -> anyhow::Result<()> { ) .await?; + // Verify task_complete notification arrives before the tool call completes. + #[expect(clippy::expect_used)] + let _task_complete = timeout( + DEFAULT_READ_TIMEOUT, + mcp_process.read_stream_until_legacy_task_complete_notification(), + ) + .await + .expect("task_complete_notification timeout") + .expect("task_complete_notification resp"); + // Verify the original `codex` tool call completes and that `git init` ran // successfully. let codex_response = timeout( diff --git a/codex-rs/mcp-server/tests/common/mcp_process.rs b/codex-rs/mcp-server/tests/common/mcp_process.rs index a659b1d9..35484264 100644 --- a/codex-rs/mcp-server/tests/common/mcp_process.rs +++ b/codex-rs/mcp-server/tests/common/mcp_process.rs @@ -474,4 +474,46 @@ impl McpProcess { })) .await } + + /// Reads notifications until a legacy TaskComplete event is observed: + /// Method "codex/event" with params.msg.type == "task_complete". + pub async fn read_stream_until_legacy_task_complete_notification( + &mut self, + ) -> anyhow::Result { + loop { + let message = self.read_jsonrpc_message().await?; + eprint!("message: {message:?}"); + + match message { + JSONRPCMessage::Notification(notification) => { + let is_match = if notification.method == "codex/event" { + if let Some(params) = ¬ification.params { + params + .get("msg") + .and_then(|m| m.get("type")) + .and_then(|t| t.as_str()) + == Some("task_complete") + } else { + false + } + } else { + false + }; + + if is_match { + return Ok(notification); + } + } + JSONRPCMessage::Request(_) => { + anyhow::bail!("unexpected JSONRPCMessage::Request: {message:?}"); + } + JSONRPCMessage::Error(_) => { + anyhow::bail!("unexpected JSONRPCMessage::Error: {message:?}"); + } + JSONRPCMessage::Response(_) => { + anyhow::bail!("unexpected JSONRPCMessage::Response: {message:?}"); + } + } + } + } }