Separate Guardian conversation preparation from execution (#55533)

## What changed

Split retained conversation scoring into `Reservation::prepare` and `PreparedConversation::run`, exposing the complete admitted input before model execution. Keep the history fork and capacity permit owned by the prepared conversation, and recheck authorization, score expiry, and backend generation before execution.

## Testing

Extend conversation tests to verify that preparation sends no model request, execution uses the prepared input, dropping a prepared conversation releases capacity, and stale prepared reviews cannot start model execution.

GitOrigin-RevId: c96cc6fac69179a78ecd7bcdb8e5a4b3087aa6e6
This commit is contained in:
swastik-oai
2026-10-20 10:52:46 +00:00
committed by copyberry
parent 46d252f8bf
commit 464add4958
6 changed files with 252 additions and 64 deletions
@@ -401,17 +401,20 @@ impl Classification {
);
}
let (ready, score) = tokio::sync::oneshot::channel();
tokio::spawn(reservation.run(ConversationRequest {
let request = ConversationRequest {
evidence,
sampling,
reset_token_limit:
guardian_config.async_classifier_conversation_token_limit,
reset_token_limit: guardian_config
.async_classifier_conversation_token_limit,
ready,
authorization: score_authorization.clone(),
thread: Arc::clone(&thread),
metrics: metrics.clone(),
pending_score: pending_score.clone(),
}));
};
if let Some(prepared) = reservation.prepare(request).await {
tokio::spawn(prepared.run());
}
tokio::select! {
biased;
() = pending_score.expired() => Err(LunaSamplerError::Superseded),
@@ -4,6 +4,7 @@
//! Permits bound preparation and sampling through full response completion.
//! Backend replacement invalidates its generation without changing snapshot shutdown policy.
//! Oversized retained requests rebuild from fresh evidence before the model's hard limit.
//! Preparation exposes the complete input without starting model execution.
//! Each request owns its fork; failed or cancelled requests leave committed history intact.
//! Host-only instruction delivery metadata is committed with its completed input messages.
//! Expired classifications release their transport without committing partial history.
@@ -23,6 +24,8 @@ use tokio::sync::Semaphore;
use tokio::sync::oneshot;
use super::authorization::ScoreAuthorization;
use super::request::PreparedRequest;
use super::sampler::ConversationSample;
use super::sampler::LunaSampler;
use super::sampler::LunaSamplerError;
use super::sampler::LunaSamplingRequest;
@@ -64,6 +67,19 @@ pub(super) struct ConversationRequest {
pub(super) pending_score: PendingScore,
}
/// Owns an admitted history fork and its permit until execution completes or it is dropped.
pub(super) struct PreparedConversation {
reservation: Reservation,
pub(super) request: PreparedRequest,
sampling: LunaSamplingRequest,
key: ReuseKey,
state: ConversationState<Vec<ResponseItemEnvelope>>,
thread: Arc<CodexThread>,
ready: oneshot::Sender<Result<SampledScore, LunaSamplerError>>,
metrics: Option<Arc<dyn ExtensionMetrics>>,
pending_score: PendingScore,
}
// Authorization covers model settings; local config changes replace the backend.
#[derive(PartialEq)]
struct ReuseKey {
@@ -109,17 +125,20 @@ impl ConversationBackend {
}
impl Reservation {
pub(super) async fn run(self, mut request: ConversationRequest) {
pub(super) async fn prepare(
self,
mut request: ConversationRequest,
) -> Option<PreparedConversation> {
request.sampling.pending_score = Some(request.pending_score.clone());
if self.generation.strong_count() == 0
|| !request.authorization.is_current(&request.thread).await
|| request.pending_score.is_expired()
{
let _ = request.ready.send(Err(LunaSamplerError::Superseded));
return;
return None;
}
let key = ReuseKey {
authorization: request.authorization.clone(),
authorization: request.authorization,
instructions: request.sampling.instructions.clone(),
parent_compaction: request.sampling.parent_compaction.clone(),
};
@@ -135,16 +154,16 @@ impl Reservation {
.map(ConversationState::from)
.unwrap_or_default();
let had_history = state.snapshot().is_some();
let prepared = request
let review_request = request
.sampling
.prepare_retained(
&request.evidence,
&mut state,
self.sampler.max_input_tokens(),
)
.filter(|prepared| {
prepared.existing_context_tokens == 0
|| prepared.input_tokens <= request.reset_token_limit
.filter(|review_request| {
review_request.existing_context_tokens == 0
|| review_request.input_tokens <= request.reset_token_limit
})
.or_else(|| {
if !had_history {
@@ -157,53 +176,83 @@ impl Reservation {
self.sampler.max_input_tokens(),
)
});
let Some(mut prepared) = prepared else {
let Some(review_request) = review_request else {
let _ = request.ready.send(Err(LunaSamplerError::InputTooLarge));
return;
return None;
};
super::metrics::record_section_costs(
request.metrics.as_deref(),
prepared.section_costs.iter().copied(),
review_request.section_costs.iter().copied(),
);
super::metrics::record_request_tokens(
request.metrics.as_deref(),
prepared.existing_context_tokens,
prepared.input_tokens,
review_request.existing_context_tokens,
review_request.input_tokens,
super::metrics::RequestAdmission::Admitted,
);
drop(request.evidence);
let cursor = prepared.cursor;
let pending_truncations = std::mem::take(&mut prepared.truncations);
Some(PreparedConversation {
reservation: self,
request: review_request,
sampling: request.sampling,
key,
state,
thread: request.thread,
ready: request.ready,
metrics: request.metrics,
pending_score: request.pending_score,
})
}
}
impl PreparedConversation {
/// Publishes the first score, then drains and commits only valid completed history.
pub(super) async fn run(mut self) {
// A caller can retain prepared input while authorization or admission expires.
if self.reservation.generation.strong_count() == 0
|| !self.key.authorization.is_current(&self.thread).await
|| self.pending_score.is_expired()
{
let _ = self.ready.send(Err(LunaSamplerError::Superseded));
return;
}
let cursor = self.request.cursor;
let pending_truncations = std::mem::take(&mut self.request.truncations);
let sample = ConversationSample {
input: self.request,
sampling: self.sampling,
};
let history = tokio::select! {
biased;
() = request.pending_score.expired() => None,
history = self.sampler.sample_retained(prepared, request.ready) => history,
() = self.pending_score.expired() => None,
history = self.reservation.sampler.sample_retained(sample, self.ready) => history,
};
if let Some(history) = history
&& self.generation.strong_count() > 0
&& request.authorization.is_current(&request.thread).await
&& self.reservation.generation.strong_count() > 0
&& self.key.authorization.is_current(&self.thread).await
{
let mut truncations = super::truncation::ClassificationTruncations::default();
truncations.extend(pending_truncations);
truncations.emit(request.metrics.as_deref());
state.complete_review(cursor);
state.commit_snapshot(history);
truncations.emit(self.metrics.as_deref());
self.state.complete_review(cursor);
self.state.commit_snapshot(history);
let mut latest = self
.reservation
.history
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if latest
.as_ref()
.is_none_or(|previous| previous.index < self.index)
.is_none_or(|previous| previous.index < self.reservation.index)
{
*latest = Some(Arc::new(CompletedConversation {
index: self.index,
key,
state,
index: self.reservation.index,
key: self.key,
state: self.state,
}));
}
}
// The permit is released only after completion and commit, not after the early score.
drop(self._permit);
drop(self.reservation._permit);
}
}
@@ -96,8 +96,21 @@ async fn configured_reset_limit_rebuilds_history_without_rejecting_fresh_evidenc
for index in 0..2 {
let (mut request, score) = pending_request(&fixture, "parent").await?;
request.metrics = Some(metrics.clone());
backend.reserve(|| index).1.unwrap().run(request).await;
let reservation = backend.reserve(|| index).1.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
let input = prepared
.request
.input
.iter()
.map(|item| item.item.clone())
.collect::<Vec<_>>();
assert_eq!(mock.requests().len(), index);
prepared.run().await;
assert_eq!(score.await??, "low");
assert_eq!(
mock.requests()[index].input(),
serde_json::to_value(input)?.as_array().unwrap().clone()
);
}
let requests = mock.requests();
assert_eq!(requests.len(), 2);
@@ -230,6 +243,16 @@ async fn delayed_preparation_does_not_block_newer_requests_and_capacity_is_bound
let mut reservations = (0..16)
.map(|index| backend.reserve(|| index).1.unwrap())
.collect::<Vec<_>>();
let (request, score) = pending_request(&fixture, "dropped").await?;
let reservation = reservations.pop().unwrap();
let prepared = reservation.prepare(request).await.unwrap();
assert!(matches!(
backend.reserve(|| 16).1,
Err(LunaSamplerError::QueueFull)
));
drop(prepared);
assert!(score.await.is_err());
reservations.push(backend.reserve(|| 15).1.unwrap());
fixture.score_tool(ToolName::plain("read_file")).await;
let cached = progress.inspect(/*call_id*/ None);
assert!(cached.has_unscored_failure);
@@ -252,11 +275,17 @@ async fn delayed_preparation_does_not_block_newer_requests_and_capacity_is_bound
// must neither read C's future evidence nor replace its retained history.
for (reservation, label) in [(c, "C"), (b, "B"), (a, "A")] {
let (request, score) = pending_request(&fixture, label).await?;
tokio::time::timeout(ASYNC_TEST_TIMEOUT, reservation.run(request)).await?;
tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
})
.await?;
assert_eq!(score.await??, "low");
}
let (request, score) = pending_request(&fixture, "D").await?;
backend.reserve(|| 3).1.unwrap().run(request).await;
let reservation = backend.reserve(|| 3).1.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
assert_eq!(score.await??, "low");
let requests = mock.requests();
let turns = requests
@@ -293,7 +322,7 @@ async fn delayed_preparation_does_not_block_newer_requests_and_capacity_is_bound
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stale_authorization_expiry_and_stopped_generations_cannot_start() -> Result<()> {
async fn stale_authorization_expiry_and_stopped_generations_reject_preparation() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let fixture = GuardianFailureFixture::with_config(
@@ -311,16 +340,20 @@ async fn stale_authorization_expiry_and_stopped_generations_cannot_start() -> Re
.insert(ConversationBackend::new(sampler));
let thread_store = fixture.test.codex.thread_extension_data();
let backend = thread_store.get::<ConversationBackend>().unwrap();
let expired = backend.reserve(|| 0).1.unwrap();
let (mut request, score) = pending_request(&fixture, "expired").await?;
request.pending_score.max_lag = 0;
observe_unscored_call(&request.pending_score.progress, thread_store);
expired.run(request).await;
let progress = Arc::clone(&request.pending_score.progress);
let (index, expired) = backend.reserve(|| observe_unscored_call(&progress, thread_store));
request.pending_score.index = index;
let expired = expired.unwrap();
let current_index = observe_unscored_call(&progress, thread_store);
assert!(expired.prepare(request).await.is_none());
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
let blocker = backend.reserve(|| 0).1.unwrap();
let stale = backend.reserve(|| 0).1.unwrap();
let (request, score) = pending_request(&fixture, "stale").await?;
let (mut request, score) = pending_request(&fixture, "stale").await?;
request.pending_score.index = current_index;
fixture
.test
.codex
@@ -328,11 +361,12 @@ async fn stale_authorization_expiry_and_stopped_generations_cannot_start() -> Re
"Stop. Do not inspect any more files.",
)])
.await?;
stale.run(request).await;
assert!(stale.prepare(request).await.is_none());
drop(blocker);
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
let stopped = backend.reserve(|| 0).1.unwrap();
let (request, score) = pending_request(&fixture, "stopped").await?;
let (mut request, score) = pending_request(&fixture, "stopped").await?;
request.pending_score.index = current_index;
drop(backend);
fixture.registry.thread_lifecycle_contributors()[0]
.on_thread_stop(ThreadStopInput {
@@ -340,7 +374,75 @@ async fn stale_authorization_expiry_and_stopped_generations_cannot_start() -> Re
thread_store,
})
.await;
stopped.run(request).await;
assert!(stopped.prepare(request).await.is_none());
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
assert!(
server
.received_requests()
.await
.unwrap()
.iter()
.all(|request| request.method.as_str() != "POST")
);
fixture.test.codex.shutdown_and_wait().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn prepared_reviews_recheck_expiry_authorization_and_stopped_generations() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let fixture = GuardianFailureFixture::with_config(
"[features.guardianv2]\nasync_classifier_mode = 'conversation'",
)
.await?;
let config = sampler_config(server.uri());
let sampler = Arc::new(LunaSampler::new(config));
fixture
.test
.codex
.thread_extension_data()
.insert(ConversationBackend::new(sampler));
let thread_store = fixture.test.codex.thread_extension_data();
let backend = thread_store.get::<ConversationBackend>().unwrap();
let (mut request, score) = pending_request(&fixture, "expired").await?;
request.pending_score.max_lag = 0;
let progress = Arc::clone(&request.pending_score.progress);
let (index, expired) = backend.reserve(|| observe_unscored_call(&progress, thread_store));
request.pending_score.index = index;
let expired = expired.unwrap();
let prepared = expired.prepare(request).await.unwrap();
let current_index = observe_unscored_call(&progress, thread_store);
prepared.run().await;
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
let blocker = backend.reserve(|| 0).1.unwrap();
let stale = backend.reserve(|| 0).1.unwrap();
let (mut request, score) = pending_request(&fixture, "stale").await?;
request.pending_score.index = current_index;
let prepared = stale.prepare(request).await.unwrap();
fixture
.test
.codex
.inject_response_items(vec![user_instruction(
"Stop. Do not inspect any more files.",
)])
.await?;
prepared.run().await;
drop(blocker);
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
let stopped = backend.reserve(|| 0).1.unwrap();
let (mut request, score) = pending_request(&fixture, "stopped").await?;
request.pending_score.index = current_index;
let prepared = stopped.prepare(request).await.unwrap();
drop(backend);
fixture.registry.thread_lifecycle_contributors()[0]
.on_thread_stop(ThreadStopInput {
session_store: &fixture.session_store,
thread_store,
})
.await;
prepared.run().await;
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
assert!(
server
@@ -404,7 +506,9 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul
let sampler = Arc::new(LunaSampler::new(config));
let backend = ConversationBackend::new(sampler);
let (request, score) = pending_request(&fixture, "seed").await?;
backend.reserve(|| 0).1.unwrap().run(request).await;
let reservation = backend.reserve(|| 0).1.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
assert_eq!(score.await??, "low");
let mut older = None;
@@ -427,7 +531,9 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul
}])
.await?;
let (request, score) = pending_request(&fixture, label).await?;
let task = tokio::spawn(backend.reserve(|| index).1.unwrap().run(request));
let reservation = backend.reserve(|| index).1.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
let task = tokio::spawn(prepared.run());
assert_eq!(
tokio::time::timeout(ASYNC_TEST_TIMEOUT, score).await???,
"low"
@@ -446,7 +552,9 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul
// A newer response with invalid retained output must not replace good history.
let (request, score) = pending_request(&fixture, "unretained").await?;
backend.reserve(|| 3).1.unwrap().run(request).await;
let reservation = backend.reserve(|| 3).1.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
assert_eq!(score.await??, "low");
if older_outcome == "completed" {
finish_older.take().unwrap().send(()).unwrap();
@@ -477,7 +585,9 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul
);
drop(finish_older);
let (request, score) = pending_request(&fixture, "probe").await?;
backend.reserve(|| 4).1.unwrap().run(request).await;
let reservation = backend.reserve(|| 4).1.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
assert_eq!(score.await??, "low");
let requests = server
@@ -585,7 +695,9 @@ async fn expiry_preserves_the_completed_checkpoint_and_transcript_cursor() -> Re
let (mut request, score) = pending_request(&fixture, "checkpoint").await?;
let (index, reservation) = backend.reserve(|| observe_unscored_call(&progress, store));
request.pending_score.index = index;
reservation.unwrap().run(request).await;
let reservation = reservation.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
assert_eq!(score.await??, "low");
// This evidence is beyond the committed cursor. A cancelled review must not
@@ -605,7 +717,9 @@ async fn expiry_preserves_the_completed_checkpoint_and_transcript_cursor() -> Re
let (mut request, early_score) = pending_request(&fixture, "expired").await?;
let (index, reservation) = backend.reserve(|| observe_unscored_call(&progress, store));
request.pending_score.index = index;
let expired = tokio::spawn(reservation.unwrap().run(request));
let reservation = reservation.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
let expired = tokio::spawn(prepared.run());
assert_eq!(early_score.await??, "low");
// Keep completion gated even after the early score, then expire the request.
for _ in 0..3 {
@@ -615,7 +729,9 @@ async fn expiry_preserves_the_completed_checkpoint_and_transcript_cursor() -> Re
let (mut request, score) = pending_request(&fixture, "resumed").await?;
let (index, reservation) = backend.reserve(|| observe_unscored_call(&progress, store));
request.pending_score.index = index;
reservation.unwrap().run(request).await;
let reservation = reservation.unwrap();
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
assert_eq!(
tokio::time::timeout(ASYNC_TEST_TIMEOUT, score).await???,
"low"
@@ -673,7 +789,12 @@ async fn expiry_discards_a_queued_score_without_recording_success() -> Result<()
request.pending_score.index = index;
let pending_score = request.pending_score.clone();
// Finish the worker while the score stays queued and the consumer remains unpolled.
tokio::time::timeout(ASYNC_TEST_TIMEOUT, reservation.unwrap().run(request)).await?;
let reservation = reservation.unwrap();
tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
let prepared = reservation.prepare(request).await.unwrap();
prepared.run().await;
})
.await?;
assert!(
!metrics
.0
@@ -22,7 +22,7 @@ impl LunaSamplingRequest {
evidence: &CollectedTranscript,
conversation: &mut ConversationState<Vec<ResponseItemEnvelope>>,
max_input_tokens: usize,
) -> Option<PreparedRequest<'_>> {
) -> Option<PreparedRequest> {
let cursor = conversation.cursor();
if conversation.snapshot().is_some() != cursor.is_some() {
return None;
@@ -68,7 +68,6 @@ impl LunaSamplingRequest {
input = history;
}
Some(PreparedRequest {
sampling: self,
input,
cursor: next_cursor,
truncations,
@@ -79,9 +78,8 @@ impl LunaSamplingRequest {
}
}
/// Admitted input and its token count, tied to the sampling settings used to prepare it.
pub(super) struct PreparedRequest<'a> {
pub(super) sampling: &'a LunaSamplingRequest,
/// Admitted input and accounting for one retained review.
pub(super) struct PreparedRequest {
pub(super) input: Vec<ResponseItemEnvelope>,
pub(super) cursor: TranscriptCursor,
pub(super) truncations: Vec<TruncationObservation>,
@@ -1,5 +1,6 @@
use super::*;
use crate::async_scorer::request::PreparedRequest;
use crate::async_scorer::sampler::ConversationSample;
use codex_guardian_context::TranscriptCursor;
use codex_history::ResponseItemEnvelope;
use pretty_assertions::assert_eq;
@@ -33,7 +34,6 @@ async fn retained_websocket_requests_preserve_model_output_and_ids() -> Result<(
.map(codex_guardian_context::estimate_input_tokens)
.fold(0usize, usize::saturating_add);
let prepared = PreparedRequest {
sampling: &request,
input: request
.input
.iter()
@@ -49,7 +49,11 @@ async fn retained_websocket_requests_preserve_model_output_and_ids() -> Result<(
truncations: Vec::new(),
section_costs: Vec::new(),
};
let (completed, score) = tokio::join!(sampler.sample_retained(prepared, ready), score);
let request = ConversationSample {
input: prepared,
sampling: request,
};
let (completed, score) = tokio::join!(sampler.sample_retained(request, ready), score);
assert_eq!(score??.accept(), "low");
let completed = completed
.expect("completed model history")
@@ -139,7 +143,6 @@ async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<(
.map(codex_guardian_context::estimate_input_tokens)
.fold(0usize, usize::saturating_add);
let prepared = PreparedRequest {
sampling: &request,
input: request
.input
.iter()
@@ -155,13 +158,17 @@ async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<(
truncations: Vec::new(),
section_costs: Vec::new(),
};
let request = ConversationSample {
input: prepared,
sampling: request,
};
let history = if receive_score {
let (history, score) = tokio::join!(sampler.sample_retained(prepared, ready), score);
let (history, score) = tokio::join!(sampler.sample_retained(request, ready), score);
assert_eq!(score??.accept(), "low");
history
} else {
let score = early_delta.then_some(score);
let history = sampler.sample_retained(prepared, ready).await;
let history = sampler.sample_retained(request, ready).await;
// A queued score discarded without expiry remains a genuine cancellation.
drop(score);
history
@@ -101,6 +101,12 @@ pub struct LunaSamplingRequest {
pub(super) pending_score: Option<PendingScore>,
}
/// Admitted retained input and Luna settings consumed together at execution.
pub(super) struct ConversationSample {
pub(super) input: PreparedRequest,
pub(super) sampling: LunaSamplingRequest,
}
/// Failures returned while connecting or sampling the Luna model.
#[derive(Debug, Error)]
pub enum LunaSamplerError {
@@ -371,9 +377,13 @@ impl LunaSampler {
/// Publishes the early score and returns only validated, completed history.
pub(super) async fn sample_retained(
&self,
prepared: PreparedRequest<'_>,
sample: ConversationSample,
ready: oneshot::Sender<Result<SampledScore, LunaSamplerError>>,
) -> Option<Vec<ResponseItemEnvelope>> {
let ConversationSample {
input: review_request,
sampling,
} = sample;
let auth_owner_generation = self.auth_owner_generation();
let mut completion = execution::RetainedCompletion {
ready: Some(ready),
@@ -383,19 +393,19 @@ impl LunaSampler {
.config
.max_input_tokens
.saturating_sub(/*rhs*/ 256)
.saturating_sub(prepared.input_tokens),
.saturating_sub(review_request.input_tokens),
};
// Transport sees model items; delivery proof follows the unchanged input prefix.
let (input, metadata): (Vec<_>, Vec<_>) = prepared
let (input, metadata): (Vec<_>, Vec<_>) = review_request
.input
.into_iter()
.map(|envelope| (envelope.item, envelope.metadata))
.unzip();
let mut execution = self.execution(prepared.sampling, input, auth_owner_generation);
let mut execution = self.execution(&sampling, input, auth_owner_generation);
execution.request.include = vec!["reasoning.encrypted_content".to_owned()];
let mut telemetry = telemetry::SamplingTelemetry::new(
self.config.metrics.clone(),
prepared.sampling.pending_score.clone(),
sampling.pending_score.clone(),
);
let result = execution
.run(