diff --git a/codex-rs/mcp-server/tests/codex_tool.rs b/codex-rs/mcp-server/tests/codex_tool.rs index 92f11eaa4f..ead8adb2a8 100644 --- a/codex-rs/mcp-server/tests/codex_tool.rs +++ b/codex-rs/mcp-server/tests/codex_tool.rs @@ -114,6 +114,13 @@ async fn shell_command_approval_triggers_elicitation() -> anyhow::Result<()> { ) .await?; + // Verify task_complete notification arrives before the tool call completes. + let _task_complete = timeout( + DEFAULT_READ_TIMEOUT, + mcp_process.read_stream_until_legacy_task_complete_notification(), + ) + .await??; + // 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 a659b1d950..35484264fa 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:?}"); + } + } + } + } }