Make the Apps recovery exposure test deterministic (#35196)

## Why

The Apps MCP background retry could complete before the test inspected the
initial sampling request, making the recovery timing nondeterministic.

## What changed

- Add a test-server gate that can hold and explicitly release the next
  successful Apps MCP initialization.
- Use a dedicated MCP mock server so the held initialization does not block
  model responses or app discovery.
- Assert that recovery uses exactly two initialization attempts and that Apps
  guidance and the deferred namespace appear only after the gate is released.

GitOrigin-RevId: deb21d4b41f353c3f84428616769af9faa1a0487
This commit is contained in:
jif
2026-07-24 14:41:21 +00:00
committed by copyberry
parent 634a998d8a
commit 5f6a2c3adb
2 changed files with 89 additions and 28 deletions
@@ -8,9 +8,11 @@ use codex_models_manager::bundled_models_response;
use serde_json::Value;
use serde_json::json;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::mpsc;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::Request;
@@ -54,6 +56,8 @@ const CALENDAR_LIST_EVENTS_RESOURCE_URI: &str = "connector://calendar/tools/cale
pub const DOCUMENT_EXTRACT_TEXT_RESOURCE_URI: &str =
"connector://calendar/tools/calendar_extract_text";
type AppsStartupInitializeGate = Arc<Mutex<Option<mpsc::Receiver<()>>>>;
#[derive(Clone)]
pub struct AppsTestServer {
pub chatgpt_base_url: String,
@@ -63,6 +67,7 @@ pub struct AppsTestServer {
pub struct AppsTestServerStartupControl {
initialize_attempts: Arc<AtomicUsize>,
remaining_initialize_failures: Arc<AtomicUsize>,
successful_initialize_gate: AppsStartupInitializeGate,
}
impl AppsTestServerStartupControl {
@@ -74,6 +79,16 @@ impl AppsTestServerStartupControl {
pub fn initialize_attempts(&self) -> usize {
self.initialize_attempts.load(Ordering::SeqCst)
}
/// Holds one successful startup without turning recovery into a fresh runtime refresh.
pub fn hold_next_successful_initialize(&self) -> mpsc::Sender<()> {
let (release, gate) = mpsc::channel();
*self
.successful_initialize_gate
.lock()
.expect("Apps initialization gate lock should not be poisoned") = Some(gate);
release
}
}
#[derive(Clone, Copy)]
@@ -177,6 +192,7 @@ impl AppsTestServer {
let control = AppsTestServerStartupControl {
initialize_attempts: Arc::new(AtomicUsize::new(0)),
remaining_initialize_failures: Arc::new(AtomicUsize::new(0)),
successful_initialize_gate: Arc::new(Mutex::new(None)),
};
mount_streamable_http_json_rpc_with_startup_control(
server,
@@ -188,6 +204,7 @@ impl AppsTestServer {
AppsTestToolsListBehavior::AlwaysAvailable,
Some(Arc::clone(&control.initialize_attempts)),
Some(Arc::clone(&control.remaining_initialize_failures)),
Some(Arc::clone(&control.successful_initialize_gate)),
)
.await;
Ok((
@@ -413,6 +430,7 @@ async fn mount_streamable_http_json_rpc_at_path(
tools_list_behavior,
/*initialize_attempts*/ None,
/*remaining_initialize_failures*/ None,
/*successful_initialize_gate*/ None,
)
.await;
}
@@ -428,6 +446,7 @@ async fn mount_streamable_http_json_rpc_with_startup_control(
tools_list_behavior: AppsTestToolsListBehavior,
initialize_attempts: Option<Arc<AtomicUsize>>,
remaining_initialize_failures: Option<Arc<AtomicUsize>>,
successful_initialize_gate: Option<AppsStartupInitializeGate>,
) {
Mock::given(method("POST"))
.and(path_regex(mcp_path_regex))
@@ -439,6 +458,7 @@ async fn mount_streamable_http_json_rpc_with_startup_control(
tools_list_behavior,
initialize_attempts,
remaining_initialize_failures,
successful_initialize_gate,
})
.mount(server)
.await;
@@ -452,6 +472,7 @@ struct CodexAppsJsonRpcResponder {
tools_list_behavior: AppsTestToolsListBehavior,
initialize_attempts: Option<Arc<AtomicUsize>>,
remaining_initialize_failures: Option<Arc<AtomicUsize>>,
successful_initialize_gate: Option<AppsStartupInitializeGate>,
}
impl Respond for CodexAppsJsonRpcResponder {
@@ -491,6 +512,20 @@ impl Respond for CodexAppsJsonRpcResponder {
"error": "simulated non-retryable Apps MCP startup failure",
}));
}
let gate = self.successful_initialize_gate.as_ref().and_then(|gate| {
gate.lock()
.expect("Apps initialization gate lock should not be poisoned")
.take()
});
// This responder's dedicated mock server stays paused while the test controls
// when this existing client's background retry may finish.
if let Some(gate) = gate
&& gate.recv().is_err()
{
return ResponseTemplate::new(500).set_body_json(json!({
"error": "Apps initialization gate was dropped before release",
}));
}
let id = body.get("id").cloned().unwrap_or(Value::Null);
let protocol_version = body
.pointer("/params/protocolVersion")
+54 -28
View File
@@ -62,6 +62,10 @@ struct CoalescingMcpContributor {
observed_markers: Mutex<Vec<String>>,
}
struct AppsMcpServerContributor {
url: String,
}
impl CoalescingMcpContributor {
fn new() -> Self {
Self {
@@ -110,6 +114,26 @@ impl McpServerContributor<Config> for CoalescingMcpContributor {
}
}
impl McpServerContributor<Config> for AppsMcpServerContributor {
fn id(&self) -> &'static str {
"deferred_apps_recovery_test"
}
fn contribute<'a>(
&'a self,
_context: McpServerContributionContext<'a, Config>,
) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
Box::pin(async move {
let config = serde_json::from_value(json!({ "url": self.url }))
.expect("test Apps MCP server config should be valid");
vec![McpServerContribution::Set {
name: CODEX_APPS_MCP_SERVER_NAME.to_string(),
config: Box::new(config),
}]
})
}
}
fn format_labeled_requests_snapshot(
scenario: &str,
sections: &[(&str, &ResponsesRequest)],
@@ -587,10 +611,17 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (apps_server, startup_control) =
AppsTestServer::mount_with_startup_control(&server).await?;
// The exact step is captured once before the first request; its background retry can recover.
// Wiremock has one server thread, so isolate the held MCP response from model/app discovery.
let apps_server_mock = responses::start_mock_server().await;
let (gated_apps_server, startup_control) =
AppsTestServer::mount_with_startup_control(&apps_server_mock).await?;
let apps_server = AppsTestServer::mount_searchable(&server).await?;
startup_control.fail_next_initialize_attempts(/*attempts*/ 1);
let release_apps_recovery = startup_control.hold_next_successful_initialize();
let mut extensions = ExtensionRegistryBuilder::new();
extensions.mcp_server_contributor(Arc::new(AppsMcpServerContributor {
url: format!("{}/api/codex/ps/mcp", gated_apps_server.chatgpt_base_url),
}));
let call_id = "pause-for-apps";
let response = mount_sse_sequence(
&server,
@@ -627,6 +658,7 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur
)
.await;
let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone())
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
config
.features
@@ -651,33 +683,13 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur
thread_settings: Default::default(),
})
.await?;
let request = tokio::time::timeout(Duration::from_secs(3), async {
let mut request = None;
let mut apps_ready = false;
while request.is_none() || !apps_ready {
let event = test
.codex
.next_event()
.await
.expect("event stream should stay open");
match event.msg {
EventMsg::RequestUserInput(next_request) => request = Some(next_request),
EventMsg::McpStartupUpdate(update)
if update.server == CODEX_APPS_MCP_SERVER_NAME
&& matches!(
update.status,
codex_protocol::protocol::McpStartupStatus::Ready
) =>
{
apps_ready = true;
}
_ => {}
}
}
request.expect("request user input event")
let EventMsg::RequestUserInput(request) = wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::RequestUserInput(_))
})
.await
.expect("Apps should recover before the second sampling request");
else {
unreachable!("wait_for_event should return the request-user-input event")
};
let initial_requests = response.requests();
assert_eq!(initial_requests.len(), 1);
@@ -700,6 +712,19 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur
"Calendar namespace should not be advertised before recovery: {initial_tools_state}"
);
release_apps_recovery
.send(())
.expect("background Apps recovery should still be waiting");
wait_for_event(&test.codex, |event| {
matches!(
event,
EventMsg::McpStartupUpdate(update)
if update.server == CODEX_APPS_MCP_SERVER_NAME
&& matches!(update.status, codex_protocol::protocol::McpStartupStatus::Ready)
)
})
.await;
test.codex
.submit(Op::UserInputAnswer {
id: request.turn_id,
@@ -757,6 +782,7 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur
}),
"recovered request should advertise tool_search: {recovered_body}"
);
assert_eq!(startup_control.initialize_attempts(), 2);
insta::assert_snapshot!(
"deferred_tools_recover_during_sampling",
format_labeled_requests_snapshot(