From cc4785e3f6e8c4121eb07ee007e5d5018b0b3f4d Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim <219906144+aibrahim-oai@users.noreply.github.com> Date: Sat, 14 Mar 2026 10:14:00 +0000 Subject: [PATCH] Buffer fuzzy search session notifications Preserve out-of-order fuzzy-search notifications in the in-process test harness so session completion events are not dropped while waiting for a later matching update. Co-authored-by: Codex --- .codex/flaky-test-triage.md | 4 + .../tests/suite/fuzzy_file_search.rs | 153 +++++++++++++----- 2 files changed, 113 insertions(+), 44 deletions(-) diff --git a/.codex/flaky-test-triage.md b/.codex/flaky-test-triage.md index dee40cdf4b..79d0c067e1 100644 --- a/.codex/flaky-test-triage.md +++ b/.codex/flaky-test-triage.md @@ -54,6 +54,7 @@ Older failures also appeared on Linux, but the repeated cross-PR signal is stron - Commit `4195e7e80` fixed that ordering race, but the new public in-process request helpers returned a private type alias and tripped the `Lint/Build` matrix before the remaining test jobs could finish. - Commit `ef5a05fa3` narrowed those signatures back to public concrete types, but the same methods still expanded to a complex nested `Result` in every `cargo clippy` target. The follow-up replaces that inline type with a documented public alias so the request-order fix does not keep failing non-release `Lint/Build`. - Commit `d87051f57` replaced the inline nested `Result` with a documented public alias, which cleared the type-complexity issue locally, but `cargo clippy` still failed in every non-release target because the new in-process fuzzy-search harness kept an unused `JSONRPCErrorError` import. The follow-up removes that stray test-only import. +- Commit `d4cbc97a9` removed that unused import, which cleared the clippy matrix and stabilized Bazel after rerunning one flaky shard, but `Tests — ubuntu-24.04-arm - aarch64-unknown-linux-gnu` still exposed an ordering race in the fuzzy-search harness: `wait_for_session_updated` discarded out-of-order `sessionCompleted` notifications, so the later completion wait could time out on faster runners. The follow-up buffers unmatched notifications and adds timeout diagnostics instead of dropping them. - Current patch set: - Pin the standalone shell test to `cmd.exe` on Windows so it validates reference-context isolation without depending on PowerShell startup behavior. - Replace the fuzzy-file-search suite's spawned `codex-app-server` harness with the in-process app-server runtime so the tests still exercise request/notification behavior without the flaky stdio startup path. @@ -61,6 +62,8 @@ Older failures also appeared on Linux, but the repeated cross-PR signal is stron - Inline the public in-process request return types so the request-order fix no longer leaks a private alias through the `app-server` public API. - Expose that request-response shape through a documented public alias so `cargo clippy -D warnings` does not reject the in-process helper API for `type_complexity`. - Remove the unused `JSONRPCErrorError` import from the in-process fuzzy-search test harness so non-release `cargo clippy --tests` can build the suite again. + - Buffer unmatched in-process notifications in the fuzzy-search harness so `sessionCompleted` events that arrive before the test starts waiting for them are preserved instead of dropped. + - Include buffered-notification method names in timeout failures so any future ordering bugs surface directly in CI annotations. - Rationale: these failures are test-harness flakes, not product behaviors. The fixes keep the assertions intact and remove environment-sensitive startup and ordering hazards instead of stretching timeouts. ## Constraints @@ -94,3 +97,4 @@ Older failures also appeared on Linux, but the repeated cross-PR signal is stron | `4195e7e80` | Preserve in-process fuzzy search request ordering | failed | Run `23084952538` proved the Linux ordering fix was necessary, but the follow-up widened the `app-server` public API with methods that returned a private type alias. `Lint/Build` failed on Linux and macOS before the remaining `Tests` jobs finished, so the next follow-up narrows the public signatures back to concrete types and reuses the same request-order behavior. | | `ef5a05fa3` | Inline in-process request response types | failed | Run `23085144063` cleared the private-alias compile break, but every non-release `Lint/Build` target still failed in `cargo clippy` while release builds passed. The next follow-up replaces the inline `IoResult>` signatures with a documented public alias, which is the smallest code change consistent with the cross-target `cargo clippy` pattern. | | `d87051f57` | Factor in-process request response alias | failed | Run `23085369065` narrowed the non-release `cargo clippy` failures down to a single test-harness warning: `unused import: codex_app_server_protocol::JSONRPCErrorError` in `app-server/tests/suite/fuzzy_file_search.rs`. The next follow-up removes that import so the request-order fix can finish the matrix. | +| `d4cbc97a9` | Remove unused fuzzy search test import | failed | Run `23085620967` cleared every completed non-release `Lint/Build` lane and passed Bazel after rerunning a flaky macOS x64 shard, but `Tests — ubuntu-24.04-arm - aarch64-unknown-linux-gnu` failed in `all::suite::fuzzy_file_search::test_fuzzy_file_search_session_multiple_query_updates_work`. The failure stack pointed to `wait_for_session_completed`, and the harness was still dropping out-of-order `sessionCompleted` notifications while waiting for matching `sessionUpdated` events. | diff --git a/codex-rs/app-server/tests/suite/fuzzy_file_search.rs b/codex-rs/app-server/tests/suite/fuzzy_file_search.rs index f144e9f372..df5f8eb338 100644 --- a/codex-rs/app-server/tests/suite/fuzzy_file_search.rs +++ b/codex-rs/app-server/tests/suite/fuzzy_file_search.rs @@ -29,6 +29,7 @@ use codex_protocol::protocol::SessionSource; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::HashMap; +use std::collections::VecDeque; use std::path::Path; use std::sync::Arc; use tempfile::TempDir; @@ -68,6 +69,7 @@ struct McpProcess { client: Option, next_request_id: i64, pending_requests: HashMap, + buffered_notifications: VecDeque, start_args: Option, } @@ -85,6 +87,7 @@ impl McpProcess { client: None, next_request_id: 1, pending_requests: HashMap::new(), + buffered_notifications: VecDeque::new(), start_args: Some(InProcessStartArgs { arg0_paths: Arg0DispatchPaths::default(), config, @@ -273,12 +276,37 @@ impl McpProcess { &mut self, method: &str, ) -> Result { + self.read_stream_until_matching_notification(method, |notification| { + Ok(notification.method == method) + }) + .await + } + + async fn read_stream_until_matching_notification

( + &mut self, + description: &str, + mut matches: P, + ) -> Result + where + P: FnMut(&JSONRPCNotification) -> Result, + { + let mut index = 0; + while index < self.buffered_notifications.len() { + if matches(&self.buffered_notifications[index])? { + if let Some(notification) = self.buffered_notifications.remove(index) { + return Ok(notification); + } + anyhow::bail!("buffered notification disappeared while waiting for {description}"); + } + index += 1; + } + loop { let event = self .client_mut()? .next_event() .await - .ok_or_else(|| anyhow!("app-server closed before emitting {method}"))?; + .ok_or_else(|| anyhow!("app-server closed before emitting {description}"))?; let notification = match event { InProcessServerEvent::ServerNotification(notification) => { @@ -286,21 +314,31 @@ impl McpProcess { } InProcessServerEvent::LegacyNotification(notification) => notification, InProcessServerEvent::Lagged { skipped } => { - anyhow::bail!("missed {skipped} app-server events while waiting for {method}") + anyhow::bail!( + "missed {skipped} app-server events while waiting for {description}" + ) } InProcessServerEvent::ServerRequest(request) => { anyhow::bail!( - "unexpected server request while waiting for {method}: {request:?}" + "unexpected server request while waiting for {description}: {request:?}" ) } }; - if notification.method == method { + if matches(¬ification)? { return Ok(notification); } + self.buffered_notifications.push_back(notification); } } + fn buffered_notification_methods(&self) -> Vec { + self.buffered_notifications + .iter() + .map(|notification| notification.method.clone()) + .collect() + } + fn client_mut(&mut self) -> Result<&mut InProcessClientHandle> { self.client .as_mut() @@ -372,54 +410,81 @@ async fn wait_for_session_updated( query: &str, file_expectation: FileExpectation, ) -> Result { - for _ in 0..20 { - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message(SESSION_UPDATED_METHOD), - ) - .await??; - let params = notification - .params - .ok_or_else(|| anyhow!("missing notification params"))?; - let payload = serde_json::from_value::(params)?; - if payload.session_id != session_id || payload.query != query { - continue; + let description = format!("session update for sessionId={session_id}, query={query}"); + let notification = match timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_matching_notification(&description, |notification| { + if notification.method != SESSION_UPDATED_METHOD { + return Ok(false); + } + let params = notification + .params + .clone() + .ok_or_else(|| anyhow!("missing notification params"))?; + let payload = + serde_json::from_value::(params)?; + let files_match = match file_expectation { + FileExpectation::Any => true, + FileExpectation::Empty => payload.files.is_empty(), + FileExpectation::NonEmpty => !payload.files.is_empty(), + }; + Ok(payload.session_id == session_id && payload.query == query && files_match) + }), + ) + .await + { + Ok(result) => result?, + Err(_) => { + anyhow::bail!( + "timed out waiting for {description}; buffered notifications={:?}", + mcp.buffered_notification_methods() + ) } - let files_match = match file_expectation { - FileExpectation::Any => true, - FileExpectation::Empty => payload.files.is_empty(), - FileExpectation::NonEmpty => !payload.files.is_empty(), - }; - if files_match { - return Ok(payload); - } - } - anyhow::bail!( - "did not receive expected session update for sessionId={session_id}, query={query}" - ); + }; + let params = notification + .params + .ok_or_else(|| anyhow!("missing notification params"))?; + Ok(serde_json::from_value::< + FuzzyFileSearchSessionUpdatedNotification, + >(params)?) } async fn wait_for_session_completed( mcp: &mut McpProcess, session_id: &str, ) -> Result { - for _ in 0..20 { - let notification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message(SESSION_COMPLETED_METHOD), - ) - .await??; - let params = notification - .params - .ok_or_else(|| anyhow!("missing notification params"))?; - let payload = - serde_json::from_value::(params)?; - if payload.session_id == session_id { - return Ok(payload); + let description = format!("session completion for sessionId={session_id}"); + let notification = match timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_matching_notification(&description, |notification| { + if notification.method != SESSION_COMPLETED_METHOD { + return Ok(false); + } + let params = notification + .params + .clone() + .ok_or_else(|| anyhow!("missing notification params"))?; + let payload = + serde_json::from_value::(params)?; + Ok(payload.session_id == session_id) + }), + ) + .await + { + Ok(result) => result?, + Err(_) => { + anyhow::bail!( + "timed out waiting for {description}; buffered notifications={:?}", + mcp.buffered_notification_methods() + ) } - } - - anyhow::bail!("did not receive expected session completion for sessionId={session_id}"); + }; + let params = notification + .params + .ok_or_else(|| anyhow!("missing notification params"))?; + Ok(serde_json::from_value::< + FuzzyFileSearchSessionCompletedNotification, + >(params)?) } async fn assert_update_request_fails_for_missing_session(