mirror of
https://github.com/openai/codex.git
synced 2026-10-29 19:08:43 +00:00
Wait for turn profiles before asserting sampling counts in trace tests (#57692)
## Why `ShutdownComplete` can arrive before detached task cleanup records and exports `codex.turn_profile` events, causing the turn phase trace test to read incomplete profile counts. ## What changed Poll exported spans until the expected number of profiles is available, with a 10-second timeout, before asserting sampling request and retry counts for reply and interrupt cases. GitOrigin-RevId: 43cdcdde71bbcd16ed2e6c8b2accaa2a5de03001
This commit is contained in:
@@ -379,26 +379,40 @@ async fn steer_joins_existing_turn_without_extending_sampling(
|
||||
assert!(dispatch_to_ready_ms.is_finite() && dispatch_to_ready_ms >= 0.0);
|
||||
assert!(accepted.end_time <= ready.timestamp);
|
||||
assert!(ready.timestamp <= blocking.end_time);
|
||||
let mut profiles: Vec<_> = spans
|
||||
.iter()
|
||||
.flat_map(|span| span.events.events.iter())
|
||||
.filter(|event| {
|
||||
event.name == "codex.turn_profile"
|
||||
&& attribute(&event.attributes, "turn_id").as_deref() == Some(call.turn_id.as_str())
|
||||
})
|
||||
.map(|event| {
|
||||
(
|
||||
attribute(&event.attributes, "conversation.id"),
|
||||
attribute(&event.attributes, "sampling_request_count"),
|
||||
attribute(&event.attributes, "sampling_retry_count"),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
profiles.sort();
|
||||
let expected_counts = match finish {
|
||||
Finish::Reply => vec!["2"],
|
||||
Finish::Interrupt => vec!["1", "1"],
|
||||
};
|
||||
// Profiles are recorded during task finalization, after run_turn ends, and
|
||||
// exported when their enclosing spans close. ShutdownComplete does not join
|
||||
// detached task cleanup, so wait for every completion before reading counts.
|
||||
let mut profiles = tokio::time::timeout(Duration::from_secs(/* secs */ 10), async {
|
||||
loop {
|
||||
let profiles: Vec<_> = exported
|
||||
.get_finished_spans()?
|
||||
.iter()
|
||||
.flat_map(|span| span.events.events.iter())
|
||||
.filter(|event| {
|
||||
event.name == "codex.turn_profile"
|
||||
&& attribute(&event.attributes, "turn_id").as_deref()
|
||||
== Some(call.turn_id.as_str())
|
||||
})
|
||||
.map(|event| {
|
||||
(
|
||||
attribute(&event.attributes, "conversation.id"),
|
||||
attribute(&event.attributes, "sampling_request_count"),
|
||||
attribute(&event.attributes, "sampling_retry_count"),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
if profiles.len() >= expected_counts.len() {
|
||||
break anyhow::Ok(profiles);
|
||||
}
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await??;
|
||||
profiles.sort();
|
||||
assert_eq!(
|
||||
profiles,
|
||||
expected_counts
|
||||
|
||||
Reference in New Issue
Block a user