Handle steers accepted during the final transcript flush (#58877)

## Why

Input accepted while a regular turn was flushing its transcript could miss the final pending-input check and remain unprocessed when the turn completed.

## What changed

Check pending input atomically with closing steer admission. If input was accepted, resume sampling in the same turn while preserving MCP startup requirements, and complete the turn with the updated answer.

## Testing

Add a regression test that pauses the final flush, submits a steer, and verifies that the model receives it exactly once before the same turn completes with the answer to the steer.

GitOrigin-RevId: 193db3ab6916a86e97c3845590598fd8fbe47081
This commit is contained in:
Charlie Marsh
2026-11-02 12:24:38 +00:00
committed by copyberry
parent 053e5b1ca9
commit ba0d0e1f74
3 changed files with 263 additions and 74 deletions
@@ -1,4 +1,4 @@
//! Exercises preparation and execution checkpoints before model requests.
//! Exercises how transcript checkpoints interact with model requests and turn completion.
use std::any::Any;
use std::sync::Arc;
@@ -49,6 +49,8 @@ use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::json;
use test_case::test_case;
use tokio::sync::Barrier;
use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::timeout;
@@ -93,6 +95,7 @@ struct GatedCheckpointStore {
arm_on: Option<PersistContext>,
checkpoint_context: Option<PersistContext>,
checkpoints: mpsc::UnboundedSender<PendingCheckpoint>,
next_flush: Mutex<Option<Arc<Barrier>>>,
}
macro_rules! delegate_store_methods {
@@ -120,10 +123,21 @@ impl ThreadStore for GatedCheckpointStore {
fn archive_thread(params: ArchiveThreadParams) -> ();
fn unarchive_thread(params: ArchiveThreadParams) -> StoredThread;
fn delete_thread(params: DeleteThreadParams) -> ();
fn flush_thread(thread_id: ThreadId) -> ();
fn shutdown_thread(thread_id: ThreadId) -> ();
}
fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {
Box::pin(async move {
let gate = self.next_flush.lock().await.take();
if let Some(gate) = gate {
// Let the test submit input while this flush is still in progress.
gate.wait().await;
gate.wait().await;
}
self.inner.flush_thread(thread_id).await
})
}
fn update_thread_metadata(
&self,
params: UpdateThreadMetadataParams,
@@ -228,6 +242,7 @@ async fn steered_input_checkpoint_controls_next_request(
arm_on: None,
checkpoint_context: None,
checkpoints,
next_flush: Mutex::default(),
});
let base_url = format!("{}/v1", server.uri());
let config_server = responses::start_mock_server().await;
@@ -348,6 +363,117 @@ async fn steered_input_checkpoint_controls_next_request(
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn steer_during_final_flush_reaches_the_model_before_turn_completion() -> anyhow::Result<()> {
let server = responses::start_mock_server().await;
let responses = responses::mount_sse_sequence(
&server,
[
("first", "original answer"),
("second", "answer to the steer"),
]
.into_iter()
.map(|(id, answer)| {
responses::sse(vec![
responses::ev_response_created(id),
responses::ev_assistant_message(id, answer),
responses::ev_completed(id),
])
})
.collect(),
)
.await;
let (checkpoints, _checkpoint_requests) = mpsc::unbounded_channel();
let store = Arc::new(GatedCheckpointStore {
inner: InMemoryThreadStore::default(),
policy: CheckpointPolicy::Synchronous,
armed: AtomicBool::new(false),
arm_on: None,
checkpoint_context: None,
checkpoints,
next_flush: Mutex::default(),
});
let test = test_codex()
.with_thread_store(store.clone())
.with_history_mode(ThreadHistoryMode::Legacy)
.build_with_auto_env(&server)
.await?;
let user_message = |text: &str, client_id| {
TurnInputRequest::new(TurnInput::UserInput {
content: vec![UserInput::Text {
text: text.to_string(),
text_elements: Vec::new(),
}],
client_id,
})
};
let flush = Arc::new(Barrier::new(/*n*/ 2));
*store.next_flush.lock().await = Some(flush.clone());
let TurnInputSubmission::Started {
turn_id,
root_turn_id,
} = test
.codex
.start_or_steer_turn(user_message("first prompt", None))
.await?
else {
panic!("expected a new turn");
};
// The regular task returned its first answer and is flushing before it completes the turn.
timeout(Duration::from_secs(10), flush.wait()).await?;
assert_eq!(responses.requests().len(), 1);
let steer = timeout(
Duration::from_secs(10),
test.codex
.start_or_steer_turn(user_message("late steer", Some("steer-id".to_string()))),
)
.await;
flush.wait().await;
assert_eq!(
steer??,
TurnInputSubmission::Steered {
turn_id: turn_id.clone(),
root_turn_id
}
);
let mut starts = 0;
let mut client_ids = Vec::new();
let EventMsg::TurnComplete(completed) = wait_for_event(&test.codex, |event| {
match event {
EventMsg::TurnStarted(_) => starts += 1,
EventMsg::UserMessage(message) if message.message == "late steer" => {
client_ids.push(message.client_id.clone());
}
_ => {}
}
matches!(event, EventMsg::TurnComplete(_))
})
.await
else {
unreachable!();
};
assert_eq!(completed.turn_id, turn_id);
assert_eq!(
completed.last_agent_message.as_deref(),
Some("answer to the steer")
);
assert_eq!(starts, 1);
assert_eq!(client_ids, [Some("steer-id".to_string())]);
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert_eq!(
requests[1]
.message_input_texts("user")
.iter()
.filter(|text| *text == "late steer")
.count(),
1
);
Ok(())
}
#[test_case(FirstInputKind::User; "user_input")]
#[test_case(FirstInputKind::InterAgentCommunication; "inter_agent_communication")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -371,6 +497,7 @@ async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata(
arm_on: None,
checkpoint_context: None,
checkpoints,
next_flush: Mutex::default(),
});
let base_url = format!("{}/v1", server.uri());
let config_server = responses::start_mock_server().await;
@@ -614,6 +741,7 @@ async fn cancelled_spawn_discards_provisional_child(caller: SpawnCaller) -> anyh
arm_on: Some(PersistContext::SubagentSpawn),
checkpoint_context: Some(PersistContext::Standard),
checkpoints,
next_flush: Mutex::default(),
});
let test = test_codex()
.with_thread_store(store)