mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user