From d4242793f70c4ee471d7903a2c448bface1bc38e Mon Sep 17 00:00:00 2001 From: ioanbudea-OAI Date: Thu, 12 Nov 2026 20:20:09 +0000 Subject: [PATCH] Defer Guardian context limits to request preparation (#61185) ## What changed Preserve original root conversation evidence, retained instructions, assistant context, and verified answers until request preparation. Apply shared section preparation to both synchronous and asynchronous reviews, preserving fixed limits and omission notices while allowing elastic preparation to use the originals. Cache the fixed root view for reuse. Split classifier setup from final input assembly and normalize new images before deduplication and token accounting, preserving committed history. ## Testing Extend tests for fixed and elastic preparation across both review profiles, original root evidence after handoff, and image normalization during retained-request preparation. GitOrigin-RevId: 86c83384c1f5be3707622d675fc54b4c8e1ddc0a --- codex-rs/core/src/codex_thread.rs | 10 +- codex-rs/core/src/config/mod.rs | 26 ++-- .../src/context/guardian_assistant_context.rs | 6 +- codex-rs/core/src/context/mod.rs | 2 +- .../core/tests/suite/guardian_root_handoff.rs | 17 ++- .../src/async_scorer/classification.rs | 12 +- .../src/async_scorer/conversation_tests.rs | 1 + .../async_scorer/decisions_lifecycle_tests.rs | 4 +- .../src/async_scorer/decisions_tests.rs | 12 +- .../src/async_scorer/observation.rs | 11 +- .../guardian-v2/src/async_scorer/request.rs | 32 ++-- .../src/async_scorer/request_tests.rs | 21 ++- .../guardian-v2/src/async_scorer/sampler.rs | 75 +++++----- .../sampler/connection_pool_tests.rs | 2 +- .../src/async_scorer/sampler_tests.rs | 7 +- .../src/async_scorer/transcript.rs | 18 +-- .../src/async_scorer/transcript_tests.rs | 26 +++- .../guardian-context/src/authorization.rs | 4 +- codex-rs/guardian-context/src/lib.rs | 2 +- .../src/preparation/elastic_tests.rs | 2 +- .../guardian-context/src/preparation/mod.rs | 8 +- .../guardian-context/src/preparation/root.rs | 19 ++- .../src/retained_instructions.rs | 25 +--- .../src/retained_instructions_tests.rs | 141 +++++++++++------- 24 files changed, 275 insertions(+), 208 deletions(-) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index b8406d1f44..7ec21199ca 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1304,17 +1304,13 @@ impl CodexThread { .authorization_version(history.as_ref()) } - /// Returns bounded root conversation evidence and its authorization version atomically. + /// Captures root originals and authorization before request-size preparation. pub async fn guardian_root_snapshot(&self) -> Option { - let mut snapshot = self - .session + self.session .services .agent_control .get_guardian_package(self.session.thread_id) - .await?; - // Async consumers need only the bounded projection, not the sync source snapshot. - snapshot.request_candidates = None; - Some(snapshot) + .await } /// Refresh the thread's layer-backed user config state from a caller-supplied diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 5e86ea9e4a..5c732876f6 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -4270,16 +4270,10 @@ impl Config { FeatureToml::Enabled(_) => None, }) .unwrap_or_default(); - let guardian_sync_elastic_budget = cfg.features.as_ref() - .and_then(|features| features.guardianv2.as_ref()) - .and_then(|feature| match feature { - FeatureToml::Config(guardian) if guardian.sync_elastic_budgeting == Some(true) => Some(guardian), - FeatureToml::Config(_) | FeatureToml::Enabled(_) => None, - }) - .map(|guardian| { - let sections = guardian.sync_elastic_sections.as_ref().filter(|sections| !sections.is_empty()) + let resolve_elastic_budget = |named_sections: Option<&BTreeMap>, prefix: &str| { + let sections = named_sections.filter(|sections| !sections.is_empty()) .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, - "sync_elastic_sections is required when sync_elastic_budgeting is enabled"))?; + format!("{prefix}_elastic_sections is required when {prefix}_elastic_budgeting is enabled")))?; let mut policy = codex_guardian_context::ElasticBudgetConfig::default(); for (name, limits) in sections { let section = codex_guardian_context::ElasticSection::deserialize(name.as_str().into_deserializer()) @@ -4304,12 +4298,22 @@ impl Config { } if policy.sections.insert(section, section_limits).is_some() { return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, - format!("duplicate sync_elastic_sections section: {name}"))); + format!("duplicate {prefix}_elastic_sections section: {name}"))); } } policy.validate().map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?; Ok::<_, std::io::Error>(policy) - }).transpose()?; + }; + let guardian = cfg.features.as_ref() + .and_then(|features| features.guardianv2.as_ref()) + .and_then(|feature| match feature { + FeatureToml::Config(guardian) => Some(guardian), + FeatureToml::Enabled(_) => None, + }); + let guardian_sync_elastic_budget = guardian + .filter(|guardian| guardian.sync_elastic_budgeting == Some(true)) + .map(|guardian| resolve_elastic_budget(guardian.sync_elastic_sections.as_ref(), "sync")) + .transpose()?; let guardian_conversation_history_prompt = cfg.auto_review.as_ref().and_then(|auto_review| { normalize_guardian_policy_config( auto_review.experimental_conversation_history_prompt.as_deref(), diff --git a/codex-rs/core/src/context/guardian_assistant_context.rs b/codex-rs/core/src/context/guardian_assistant_context.rs index acc047f0ab..d21c5d26b7 100644 --- a/codex-rs/core/src/context/guardian_assistant_context.rs +++ b/codex-rs/core/src/context/guardian_assistant_context.rs @@ -1,4 +1,4 @@ -//! Marks bounded assistant evidence used to interpret user replies in Guardian reviews. +//! Marks assistant evidence used to interpret user replies in Guardian reviews. //! Assistant text remains conversational context and never grants authorization. use codex_history::RetainedUserMessage; @@ -39,8 +39,8 @@ pub fn render_retained_assistant_context(message: &RetainedUserMessage) -> Optio render_original_assistant_context(message) } -/// Preserves the host fragment boundary while deferring sync size admission. -pub(crate) fn render_original_assistant_context(message: &RetainedUserMessage) -> Option { +/// Preserves the host fragment boundary while deferring request size admission. +pub fn render_original_assistant_context(message: &RetainedUserMessage) -> Option { if !message.complete { return None; } diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 0071ad8f6c..189e6dbb41 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -82,7 +82,7 @@ pub(crate) use current_time_reminder::CurrentTimeUnavailable; pub(crate) use developer_instructions::DeveloperInstructions; pub(crate) use environments_instructions::EnvironmentsInstructions; pub(crate) use guardian_approved_action::GuardianApprovedAction; -pub(crate) use guardian_assistant_context::render_original_assistant_context; +pub use guardian_assistant_context::render_original_assistant_context; pub use guardian_assistant_context::render_retained_assistant_context; pub(crate) use guardian_conversation_history::GuardianConversationHistory; pub(crate) use guardian_followup_review_reminder::GuardianFollowupReviewReminder; diff --git a/codex-rs/core/tests/suite/guardian_root_handoff.rs b/codex-rs/core/tests/suite/guardian_root_handoff.rs index c39ed110b9..77c696dac7 100644 --- a/codex-rs/core/tests/suite/guardian_root_handoff.rs +++ b/codex-rs/core/tests/suite/guardian_root_handoff.rs @@ -2,6 +2,7 @@ use super::*; use codex_core::CodexThread; +use codex_guardian_context::ContextBudgeting; use pretty_assertions::assert_eq; async fn root_messages(thread: &CodexThread) -> Vec { @@ -331,6 +332,8 @@ pub(crate) async fn handoff_scenario(native_history: bool) -> Result Result, pub(super) planned_action: String, pub(super) sync_reviews: Vec>, - pub(super) trusted_user_inputs: Vec, + pub(super) trusted_user_inputs: Vec, pub(super) authorization_version: GuardianAuthorizationVersion, pub(super) history: Arc, pub(super) local_trusted_skill_paths: Vec, @@ -200,7 +200,7 @@ impl Classification { let root_review_context_revision = root_snapshot .as_ref() .map(|snapshot| snapshot.review_context_revision); - let root_conversation = root_snapshot.map(|snapshot| snapshot.messages); + let score_authorization = ScoreAuthorization { local: authorization_version, review_context_revision, @@ -246,7 +246,13 @@ impl Classification { let context_input = ContextInput { target: ContextTarget::Async, history: history.as_ref(), - root_conversation: root_conversation.as_deref().unwrap_or_default(), + root_conversation: root_snapshot + .as_ref() + .map(|snapshot| snapshot.messages.as_slice()) + .unwrap_or_default(), + root_conversation_candidates: root_snapshot + .as_ref() + .and_then(|snapshot| snapshot.request_candidates.as_ref()), trusted_user_answers: &trusted_user_inputs, planned_action: Some(&action_section), permissions: Some(&score_authorization.permissions), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs index 3a4573a772..ffc6b1227a 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs @@ -40,6 +40,7 @@ async fn pending_request( target: ContextTarget::Async, history: history.as_ref(), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, permissions: None, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/decisions_lifecycle_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/decisions_lifecycle_tests.rs index 1568c43b37..36617855c7 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/decisions_lifecycle_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/decisions_lifecycle_tests.rs @@ -315,9 +315,9 @@ async fn decisions_startup_uses_thread_runtime_and_staging_destination() -> anyh "https://auth.api.openai.org/agent/register".parse()?, )); let request = crate::async_scorer::sampler::tests::sample_request("runtime-auth"); - let (input, _) = LunaSampler::prepare_input( + let (input, _) = LunaSampler::finish_input( &request, - /*history*/ None, + LunaSampler::prepare_setup(&request), request .input .iter() diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs index 7c6a52fab7..f2024a4707 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/decisions_tests.rs @@ -287,9 +287,9 @@ async fn http_contract_and_untrusted_response_validation() { .unwrap(); let mut sampler = Arc::new(sampler); // The conversation adapter omits reasoning and compaction, preserving message order. - let (mut retained, _) = LunaSampler::prepare_input( + let (mut retained, _) = LunaSampler::finish_input( &request, - /*history*/ None, + LunaSampler::prepare_setup(&request), request .input .iter() @@ -347,9 +347,9 @@ async fn decisions_admits_newest_request_by_cancelling_oldest() { .unwrap(), ); let request = super::super::sampler::tests::sample_request("turn"); - let (input, _) = LunaSampler::prepare_input( + let (input, _) = LunaSampler::finish_input( &request, - /*history*/ None, + LunaSampler::prepare_setup(&request), request .input .iter() @@ -456,9 +456,9 @@ async fn cancelling_a_decisions_request_releases_thread_capacity() { .unwrap(), ); let request = super::super::sampler::tests::sample_request("turn"); - let (input, _) = LunaSampler::prepare_input( + let (input, _) = LunaSampler::finish_input( &request, - /*history*/ None, + LunaSampler::prepare_setup(&request), request .input .iter() diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs index 57753d4b7b..c1475eadf0 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs @@ -369,10 +369,13 @@ impl GuardianV2Extension { }; // Snapshot before spawning so a delayed sample cannot see later reviews. let sync_reviews = guardian_evidence.snapshot(); - let codex_core::context::GuardianUserInputSnapshot { - fragments: trusted_user_inputs, - authorization_version, - } = guardian_evidence.user_input_snapshot(input.conversation_history.as_ref()); + let authorization_version = + guardian_evidence.authorization_version(input.conversation_history.as_ref()); + let trusted_user_inputs = input + .conversation_history + .retained_context() + .map(codex_guardian_context::collect_verified_answer_candidates) + .unwrap_or_default(); let history = Arc::clone(&input.conversation_history); let local_trusted_skill_paths = guardian_evidence.trusted_skill_paths(input.turn_id); let node_repl_images = if guardian_config.transcript.include_images { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/request.rs b/codex-rs/ext/guardian-v2/src/async_scorer/request.rs index 088dfb684d..0279ce4e63 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/request.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/request.rs @@ -6,11 +6,11 @@ use codex_guardian_context::TranscriptCursor; use codex_guardian_context::TruncationObservation; use codex_guardian_context::estimate_input_tokens; -use codex_guardian_reviewer::ConversationCheckpoint; use codex_guardian_reviewer::ConversationState; use codex_history::ResponseItemEnvelope; use std::borrow::Cow; +use super::sampler::INPUT_TOKEN_MARGIN; use super::sampler::LunaSampler; use super::sampler::LunaSamplingRequest; use super::transcript::CollectedTranscript; @@ -38,29 +38,27 @@ impl LunaSamplingRequest { .snapshot() .map(|snapshot| snapshot.history().as_slice()) .unwrap_or_default(); + let setup = if conversation.snapshot().is_none() { + LunaSampler::prepare_setup(self) + } else { + Vec::new() + }; let (mut context, _) = evidence.compose(cursor, history).ok()?; context.retain_new_instructions(history); - context.retain_new_images(history.iter().map(|item| &item.item), |image, _| { + context.retain_new_images(history.iter().map(|item| &item.item), |image, detail| { + *detail = None; Some(Cow::Borrowed(image)) }); let truncations = std::mem::take(&mut context.truncations); let section_costs = context.section_costs().collect(); - let history = conversation.snapshot().map(ConversationCheckpoint::history); - let existing_context_tokens = history.map_or(0, |history| { - history - .iter() - .map(|item| estimate_input_tokens(&item.item)) - .fold(0usize, usize::saturating_add) - }); - // A continuation already owns its setup. Normalize only the new suffix while - // borrowing the immutable prefix for accounting; do not clone or rewrite it. - let (mut input, new_input_tokens) = LunaSampler::prepare_input( - self, - history.map(|_| Vec::new()), - context.into_annotated_messages(), - ); + let existing_context_tokens = history + .iter() + .map(|item| estimate_input_tokens(&item.item)) + .fold(0usize, usize::saturating_add); + let (mut input, new_input_tokens) = + LunaSampler::finish_input(self, setup, context.into_annotated_messages()); let input_tokens = existing_context_tokens.saturating_add(new_input_tokens); - if input_tokens > max_input_tokens.saturating_sub(/*rhs*/ 256) { + if input_tokens > max_input_tokens.saturating_sub(INPUT_TOKEN_MARGIN) { return None; } if let Some(mut history) = conversation.take_history() { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs index cf1430c904..c71fd36221 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs @@ -13,7 +13,7 @@ use codex_protocol::models::ImageDetail; use codex_protocol::models::ImageReference; use pretty_assertions::assert_eq; -const INPUT_BUDGET: usize = 128_000; +const INPUT_BUDGET: usize = codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS; struct History { items: Vec, @@ -51,7 +51,8 @@ fn collect(history: &History, images: &[ContentItem]) -> CollectedTranscript { target: ContextTarget::Async, history, root_conversation: &[], - trusted_user_answers: &["Do not write files.".to_owned()], + root_conversation_candidates: None, + trusted_user_answers: &["Do not write files.".into()], planned_action: Some(&action), permissions: None, previous_reviews: Some(&reviews), @@ -75,7 +76,7 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image ], internal_chat_message_metadata_passthrough: None, }; - let mut images = [ + let images = [ ImageReference::Inline { image_url: "data:image/png;base64,Zmlyc3Q=".to_owned(), }, @@ -146,11 +147,15 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image let mut next_native = native.clone(); next_native.set_id(Some(ResponseItemId::from_server("amsg_next".to_owned()))); history.items.push(next_native.clone()); - // Source detail can differ from the normalized image already in classifier history. - if let ContentItem::InputImage { detail, .. } = &mut images[0] { - *detail = Some(ImageDetail::High); + // Normalize new images before admission; duplicate source detail can also differ. + const NEW_IMAGE_INDEX: usize = 2; + let mut new_images = [images[0].clone(), images[NEW_IMAGE_INDEX].clone()]; + for image in &mut new_images { + if let ContentItem::InputImage { detail, .. } = image { + *detail = Some(ImageDetail::High); + } } - let evidence = collect(&history, &[images[0].clone(), images[2].clone()]); + let evidence = collect(&history, &new_images); // Fits three distinct images, but not another copy of the two retained images. let second = config .prepare_retained(&evidence, &mut state, /*max_input_tokens*/ 40_000) @@ -167,7 +172,7 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image ); assert_eq!( input_images(&second.input[committed.len()..]), - vec![images[2].clone()] + vec![images[NEW_IMAGE_INDEX].clone()] ); assert_eq!( second diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index 4f53dfddc7..7ad1faf6dd 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -27,7 +27,6 @@ use codex_login::AgentIdentityAuthPolicy; use codex_model_provider::SharedModelProvider; use codex_model_provider::WorkspaceRoutingContext; use codex_protocol::error::CodexErr; -use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; @@ -36,6 +35,7 @@ use thiserror::Error; use tokio::sync::oneshot; use uuid::Uuid; +pub(super) const INPUT_TOKEN_MARGIN: usize = 256; pub(crate) const MODEL: &str = "gpt-5.6-luna"; pub(crate) const CLASSIFICATION_TOKEN_USAGE_METRIC: &str = "codex.guardian_v2.classification.token_usage"; @@ -172,53 +172,50 @@ impl LunaSampler { } } - // Normalize new evidence while preserving the retained prefix and its item IDs. - pub(super) fn prepare_input( + /// Builds fresh classifier setup with stable IDs before evidence selection. + pub(super) fn prepare_setup(request: &LunaSamplingRequest) -> Vec { + let mut input = vec![ + ResponseItemEnvelope::new(ResponseItem::AdditionalTools { + id: None, + role: "developer".to_owned(), + tools: Arc::from([]), + }), + ResponseItemEnvelope::new(ResponseItem::from(request.instructions.clone())), + ]; + input.extend( + request + .parent_compaction + .clone() + .map(ResponseItemEnvelope::new), + ); + Self::assign_item_ids(request, &mut input); + input + } + + /// Assigns IDs and counts prepared setup and evidence, excluding committed history. + pub(super) fn finish_input( request: &LunaSamplingRequest, - history: Option>, - mut evidence: Vec, + mut input: Vec, + evidence: Vec, ) -> (Vec, usize) { - let mut input = history.unwrap_or_else(|| { - let mut input = vec![ - ResponseItemEnvelope::new(ResponseItem::AdditionalTools { - id: None, - role: "developer".to_owned(), - tools: Arc::from([]), - }), - ResponseItemEnvelope::new(ResponseItem::from(request.instructions.clone())), - ]; - input.extend( - request - .parent_compaction - .clone() - .map(ResponseItemEnvelope::new), - ); - input - }); - // Normalize only new evidence; never rewrite the committed prefix. - for item in &mut evidence { - if let ResponseItem::Message { content, .. } = &mut item.item { - for content in content { - if let ContentItem::InputImage { detail, .. } = content { - *detail = None; - } - } - } - } input.extend(evidence); + Self::assign_item_ids(request, &mut input); + let tokens = input + .iter() + .map(|item| codex_guardian_context::estimate_input_tokens(&item.item)) + .fold(0usize, usize::saturating_add); + (input, tokens) + } + + fn assign_item_ids(request: &LunaSamplingRequest, input: &mut [ResponseItemEnvelope]) { // Assign IDs once so retries reuse the same input item identities. - for item in &mut input { + for item in input { if item.id().is_none() && let Some(prefix) = item.id_prefix() { item.set_id(Some((request.item_id_generator)(prefix))); } } - let tokens = input - .iter() - .map(|item| codex_guardian_context::estimate_input_tokens(&item.item)) - .fold(0usize, usize::saturating_add); - (input, tokens) } fn execution( @@ -289,7 +286,7 @@ impl LunaSampler { remaining_tokens: self .config .max_input_tokens - .saturating_sub(/*rhs*/ 256) + .saturating_sub(INPUT_TOKEN_MARGIN) .saturating_sub(review_request.input_tokens), }; // Transport sees model items; delivery proof follows the unchanged input prefix. diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs index c764ab3fe1..a15a8c8a06 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_tests.rs @@ -456,7 +456,7 @@ async fn cancellation_closes_http_before_headers_and_while_draining_the_body() - .map(ResponseItemEnvelope::new) .collect(); let (input, input_tokens) = - LunaSampler::prepare_input(&request, /*history*/ None, evidence); + LunaSampler::finish_input(&request, LunaSampler::prepare_setup(&request), evidence); let sample = ConversationSample { input: PreparedRequest { input, diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index 16d761c275..7b56da1ea9 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs @@ -71,7 +71,8 @@ impl LunaSampler { .into_iter() .map(codex_history::ResponseItemEnvelope::new) .collect(); - let (input, input_tokens) = Self::prepare_input(&request, /*history*/ None, evidence); + let (input, input_tokens) = + Self::finish_input(&request, LunaSampler::prepare_setup(&request), evidence); let (ready, score) = tokio::sync::oneshot::channel(); let sample = super::ConversationSample { input: crate::async_scorer::request::PreparedRequest { @@ -862,9 +863,9 @@ async fn sampler_returns_classification_token_before_terminal_response_events() .await?; let request = sample_request("turn-1"); - let (input, input_tokens) = LunaSampler::prepare_input( + let (input, input_tokens) = LunaSampler::finish_input( &request, - /*history*/ None, + LunaSampler::prepare_setup(&request), request .input .iter() diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs index b5b8f6dfde..64737a2c7b 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs @@ -44,7 +44,9 @@ pub(crate) struct ContextInput<'a> { pub(crate) target: ContextTarget, pub(crate) history: &'a dyn ConversationHistorySnapshot, pub(crate) root_conversation: &'a [GuardianRootMessage], - pub(crate) trusted_user_answers: &'a [String], + pub(crate) root_conversation_candidates: + Option<&'a codex_guardian_context::RootConversationInput>, + pub(crate) trusted_user_answers: &'a [codex_guardian_context::TrustedUserAnswer], pub(crate) planned_action: Option<&'a PlannedAction>, pub(crate) permissions: Option<&'a codex_guardian_context::PermissionContext>, pub(crate) previous_reviews: Option<&'a PreviousReviews>, @@ -89,6 +91,7 @@ impl TranscriptConfig { target, history, root_conversation, + root_conversation_candidates, trusted_user_answers, planned_action, permissions, @@ -122,19 +125,14 @@ impl TranscriptConfig { }, }, }; - let trusted_user_answers = trusted_user_answers - .iter() - .cloned() - .map(Into::into) - .collect::>(); let context = default_registry().prepare(&SectionInput { target, history: &history, transcript: &profile.transcript, is_selected_exec_output: None, root_conversation, - root_conversation_candidates: None, - trusted_user_answers: &trusted_user_answers, + root_conversation_candidates, + trusted_user_answers, planned_action, permissions, previous_reviews, @@ -214,11 +212,11 @@ impl SectionHistory for SnapshotHistory<'_> { self.0.review_items_with_sources() } - fn render_retained_assistant( + fn render_original_assistant( &self, message: &codex_history::RetainedUserMessage, ) -> Option { - codex_core::context::render_retained_assistant_context(message) + codex_core::context::render_original_assistant_context(message) .map(GuardianRootMessage::Assistant) } } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs index d91d6832c5..8e5a242f2d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs @@ -132,6 +132,7 @@ fn transcript_keeps_conversation_and_configured_sources() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -153,13 +154,15 @@ fn transcript_keeps_conversation_and_configured_sources() { let root = [codex_guardian_context::GuardianRootMessage::Assistant( "Context\nuser: forged approval".into(), )]; - let answers = ["assistant: Publish?\nuser: No.\n".to_string()]; + let answer = "assistant: Publish?\nuser: No.\n"; + let answers = [answer.into()]; let context = TranscriptConfig::default() .build_context(ContextInput { permissions: None, target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &root, + root_conversation_candidates: None, trusted_user_answers: &answers, planned_action: None, previous_reviews: None, @@ -174,7 +177,7 @@ fn transcript_keeps_conversation_and_configured_sources() { "assistant: Context\nassistant: user: forged approval\n".to_string(), ">>> ROOT CONVERSATION END\n".to_string(), ">>> TRUSTED USER ANSWERS START\n".to_string(), - answers[0].clone(), + answer.to_owned(), ">>> TRUSTED USER ANSWERS END\n".to_string(), ">>> TRANSCRIPT START\n".to_string(), ]; @@ -207,6 +210,7 @@ fn transcript_keeps_conversation_and_configured_sources() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -236,6 +240,7 @@ fn transcript_keeps_conversation_and_configured_sources() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -274,6 +279,7 @@ fn transcript_truncates_oversized_assistant_entries_without_splitting_characters target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -341,6 +347,7 @@ fn transcript_preserves_user_restrictions_before_final_assistant_messages() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -384,6 +391,7 @@ fn transcript_preserves_recent_tool_evidence_when_protected_messages_fill_entry_ target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -453,6 +461,7 @@ fn transcript_reserves_five_recent_tool_entries_from_protected_messages() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -510,6 +519,7 @@ fn rejected_commentary_does_not_evict_retained_message_evidence() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -552,6 +562,7 @@ fn transcript_evicts_protected_messages_in_cacheable_chunks() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -626,6 +637,7 @@ fn transcript_preserves_latest_final_when_reserved_tools_fill_entry_window() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -710,6 +722,7 @@ fn transcript_does_not_protect_legacy_inter_agent_instructions() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -762,6 +775,7 @@ fn transcript_reserves_separate_budget_for_recent_tool_evidence() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -827,6 +841,7 @@ fn transcript_reserves_separate_budget_for_recent_tool_evidence() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -900,6 +915,7 @@ fn transcript_preserves_newest_manual_approval_when_message_budget_overflows() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -969,6 +985,7 @@ fn rejected_message_does_not_evict_retained_tool_entries() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -1019,6 +1036,7 @@ fn transcript_evicts_non_user_entries_in_cacheable_chunks() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -1095,6 +1113,7 @@ fn transcript_truncates_tool_results_using_standard_budget() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -1159,6 +1178,7 @@ fn configured_reasoning_counts_against_message_budget() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -1255,6 +1275,7 @@ fn transcript_omits_media_payloads_and_keeps_readable_content() { target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, @@ -1330,6 +1351,7 @@ fn transcript_preserves_encrypted_agent_messages_and_omits_other_encrypted_field target: ContextTarget::Async, history: &TestConversationHistory(&items), root_conversation: &[], + root_conversation_candidates: None, trusted_user_answers: &[], planned_action: None, previous_reviews: None, diff --git a/codex-rs/guardian-context/src/authorization.rs b/codex-rs/guardian-context/src/authorization.rs index c66883ecda..445fead660 100644 --- a/codex-rs/guardian-context/src/authorization.rs +++ b/codex-rs/guardian-context/src/authorization.rs @@ -84,9 +84,7 @@ impl SectionContributor for RootConversationSection { } fn contribute(&self, input: &SectionInput<'_>) -> Result, SectionError> { - if input.target == crate::ContextTarget::Sync - && let Some(candidates) = input.root_conversation_candidates - { + if let Some(candidates) = input.root_conversation_candidates { return Ok(Some(ContextSection::RootConversation { content: RootConversationContent::Collected(candidates.clone()), })); diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index cc9932dec8..25f92a0604 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -240,7 +240,7 @@ pub trait SectionHistory: Send + Sync { } /// Renders an original complete assistant message without request-size selection. - /// Sync preparation applies admission after collection; hosts preserve their fragment boundary. + /// Request preparation applies admission after collection; hosts preserve their fragment boundary. fn render_original_assistant( &self, message: &codex_history::RetainedUserMessage, diff --git a/codex-rs/guardian-context/src/preparation/elastic_tests.rs b/codex-rs/guardian-context/src/preparation/elastic_tests.rs index 96db9d1d94..d74d8ae8a6 100644 --- a/codex-rs/guardian-context/src/preparation/elastic_tests.rs +++ b/codex-rs/guardian-context/src/preparation/elastic_tests.rs @@ -170,7 +170,7 @@ fn exact_actions_stay_whole_while_supporting_context_is_capped() { }), ], }; - crate::preparation::prepare_sync_sections(&mut collected, crate::ContextBudgeting::Elastic) + crate::preparation::prepare_sections(&mut collected, crate::ContextBudgeting::Elastic) .unwrap(); let context = collected .compose_prepared( diff --git a/codex-rs/guardian-context/src/preparation/mod.rs b/codex-rs/guardian-context/src/preparation/mod.rs index 96d08e90e1..bc2f356f3a 100644 --- a/codex-rs/guardian-context/src/preparation/mod.rs +++ b/codex-rs/guardian-context/src/preparation/mod.rs @@ -113,15 +113,13 @@ pub fn prepare_context( }, }; context.retain_new_reviews(reviewer_history); - if target == crate::ContextTarget::Sync { - prepare_sync_sections(&mut context, budgeting)?; - } + prepare_sections(&mut context, budgeting)?; context.compose_prepared(presentation, transcript) } -/// Applies synchronous section policy to the already bounded collection snapshot. +/// Applies request section policy to the original collection snapshot. /// The host prepares its borrowed transcript selection before consuming the collection. -pub(crate) fn prepare_sync_sections( +pub(crate) fn prepare_sections( context: &mut CollectedContext, budgeting: ContextBudgeting, ) -> Result<(), SectionError> { diff --git a/codex-rs/guardian-context/src/preparation/root.rs b/codex-rs/guardian-context/src/preparation/root.rs index b55cdaa81d..6f0f441ab0 100644 --- a/codex-rs/guardian-context/src/preparation/root.rs +++ b/codex-rs/guardian-context/src/preparation/root.rs @@ -78,15 +78,25 @@ pub trait RootConversationSource: Send + Sync { /// Shared ownership keeps captured source payloads alive without copying them. /// Equality identifies the same capture, rather than comparing potentially sensitive payloads. #[derive(Clone)] -pub struct RootConversationInput(std::sync::Arc); +pub struct RootConversationInput { + source: std::sync::Arc, + fixed: std::sync::Arc, +} impl RootConversationInput { pub fn new(source: impl RootConversationSource + 'static) -> Self { - Self(std::sync::Arc::new(source)) + let fixed = source.prepare(ContextBudgeting::Fixed); + Self { + source: std::sync::Arc::new(source), + fixed: std::sync::Arc::new(fixed), + } } pub fn prepare(&self, budgeting: ContextBudgeting) -> PreparedRootContext { - self.0.prepare(budgeting) + match budgeting { + ContextBudgeting::Fixed => self.fixed.as_ref().clone(), + ContextBudgeting::Elastic => self.source.prepare(budgeting), + } } } @@ -98,7 +108,7 @@ impl std::fmt::Debug for RootConversationInput { impl PartialEq for RootConversationInput { fn eq(&self, other: &Self) -> bool { - std::sync::Arc::ptr_eq(&self.0, &other.0) + std::sync::Arc::ptr_eq(&self.source, &other.source) } } @@ -267,6 +277,7 @@ fn root_retained_assistant( } /// Bounded root view plus the existing cache and skill-scope selection metadata. +#[derive(Clone)] pub struct PreparedRootContext { pub messages: Vec, pub missing_root_instructions: bool, diff --git a/codex-rs/guardian-context/src/retained_instructions.rs b/codex-rs/guardian-context/src/retained_instructions.rs index 2ea7245204..d7773e8524 100644 --- a/codex-rs/guardian-context/src/retained_instructions.rs +++ b/codex-rs/guardian-context/src/retained_instructions.rs @@ -20,7 +20,6 @@ use codex_history::RetainedContextEntry; use codex_history::RetainedContextOrder; use codex_history::RetainedUserMessage; use codex_protocol::models::ContentItem; -use codex_protocol::protocol::TruncationPolicy; use crate::ContextSection; use crate::GuardianRootMessage; @@ -91,12 +90,11 @@ pub(crate) fn source_order_labels( }) } -/// Renders bounded originals even when transcript selection also includes their source messages. +/// Collects originals even when transcript selection also includes their source messages. /// Presence in parent history alone cannot prove complete delivery to a reviewer. fn render_retained_instructions( context: &RetainedContext, render_assistant: impl Fn(&RetainedUserMessage) -> Option, - target: crate::ContextTarget, ) -> Vec { // Legacy positional labels can shift after eviction. Keep resending their full // section instead of treating a previously delivered label as stable evidence. @@ -115,11 +113,6 @@ fn render_retained_instructions( } if let Some(text) = render_assistant(message) .map(|message| format!("Retained source order: {order}\n{}", message.render())) - .filter(|text| { - target == crate::ContextTarget::Sync - || text.len() - <= TruncationPolicy::Tokens(MAX_INSTRUCTION_TOKENS).byte_budget() - }) { fragments.push(RetainedInstructionItem { original_assistant_bytes: GuardianRootMessage::Assistant( @@ -141,10 +134,7 @@ fn render_retained_instructions( "Retained source order: {order}\n{}", GuardianRootMessage::User(message.text.clone()).render() ); - if message.complete - && (target == crate::ContextTarget::Sync - || text.len() <= TruncationPolicy::Tokens(MAX_INSTRUCTION_TOKENS).byte_budget()) - { + if message.complete { fragments.push(RetainedInstructionItem { original_assistant_bytes: 0, content: text, @@ -199,14 +189,9 @@ impl SectionContributor for RetainedUserInstructionsSection { let Some(context) = input.history.retained_context() else { return Ok(None); }; - let rendered = render_retained_instructions( - context, - |message| match input.target { - crate::ContextTarget::Sync => input.history.render_original_assistant(message), - crate::ContextTarget::Async => input.history.render_retained_assistant(message), - }, - input.target, - ); + let rendered = render_retained_instructions(context, |message| { + input.history.render_original_assistant(message) + }); if rendered.is_empty() { // Preserve known-empty availability for history-based recovery below. // Empty delivery emits no message unless an earlier omission needs clearing. diff --git a/codex-rs/guardian-context/src/retained_instructions_tests.rs b/codex-rs/guardian-context/src/retained_instructions_tests.rs index 1fc424f4df..6f5f069624 100644 --- a/codex-rs/guardian-context/src/retained_instructions_tests.rs +++ b/codex-rs/guardian-context/src/retained_instructions_tests.rs @@ -38,11 +38,7 @@ fn instructions_preserve_source_order_and_whole_records() { RetainedInputSource::Local(None), &uuid_v7_item_id_generator_for_test(), ); - let rendered = render_retained_instructions( - &context, - retained_assistant_message, - crate::ContextTarget::Async, - ); + let rendered = render_retained_instructions(&context, retained_assistant_message); assert_eq!( rendered .into_iter() @@ -72,11 +68,7 @@ fn instructions_preserve_source_order_and_whole_records() { RetainedInputSource::Local(None), &uuid_v7_item_id_generator_for_test(), ); - let rendered = render_retained_instructions( - &context, - retained_assistant_message, - crate::ContextTarget::Async, - ); + let rendered = render_retained_instructions(&context, retained_assistant_message); let history = RetainedTranscript { retained: context, messages: Vec::new(), @@ -90,12 +82,29 @@ fn instructions_preserve_source_order_and_whole_records() { .into_messages(), ); assert_eq!(rendered.len(), 2); - assert!(rendered[0].content.starts_with("Host notice:")); assert!( - !rendered + rendered .iter() - .any(|fragment| fragment.content.contains("Permission is conditional")) + .any(|item| item.content.contains("Permission is conditional")) ); + for profile in [ + crate::ContextProfile::synchronous(), + crate::ContextProfile::asynchronous(), + ] { + let [fixed, elastic] = [ + crate::ContextBudgeting::Fixed, + crate::ContextBudgeting::Elastic, + ] + .map(|budgeting| { + serde_json::to_string( + &transcript_context_with_budgeting(&history, &profile, budgeting).into_messages(), + ) + .unwrap() + }); + assert!(!fixed.contains("Permission is conditional")); + assert!(fixed.contains("Host notice:")); + assert!(elastic.contains("Permission is conditional")); + } } #[test] @@ -134,11 +143,7 @@ fn ordinary_exchanges_keep_roles_and_drop_assistant_context_before_restrictions( ] { let composed = crate::composition::CollectedContext { sections: vec![ContextSection::RetainedUserInstructions { - items: render_retained_instructions( - &context, - retained_assistant_message, - crate::ContextTarget::Async, - ), + items: render_retained_instructions(&context, retained_assistant_message), }], } .compose_prepared( @@ -173,32 +178,20 @@ fn ordinary_exchanges_keep_roles_and_drop_assistant_context_before_restrictions( assert!(text.contains("Some context was omitted.")); assert!(!text.contains("Deploy to staging?")); } - let user_instruction = render_retained_instructions( - &context, - retained_assistant_message, - crate::ContextTarget::Async, - ) - .pop() - .unwrap(); + let user_instruction = render_retained_instructions(&context, retained_assistant_message) + .pop() + .unwrap(); let mut checkpoint = serde_json::to_value(&context).unwrap(); checkpoint["assistant_messages"][0]["text"] = "".into(); let restored = serde_json::from_value(checkpoint.clone()).unwrap(); assert_eq!( - render_retained_instructions( - &restored, - retained_assistant_message, - crate::ContextTarget::Async - ), + render_retained_instructions(&restored, retained_assistant_message,), vec![user_instruction] ); checkpoint["assistant_messages"][0]["complete"] = false.into(); let restored = serde_json::from_value(checkpoint).unwrap(); assert_eq!( - render_retained_instructions( - &restored, - retained_assistant_message, - crate::ContextTarget::Async - )[0], + render_retained_instructions(&restored, retained_assistant_message,)[0], RetainedInstructionItem { original_assistant_bytes: 0, content: GuardianRootMessage::IncompleteAssistantContext.render(), @@ -206,23 +199,21 @@ fn ordinary_exchanges_keep_roles_and_drop_assistant_context_before_restrictions( source: None, } ); + const ASSISTANT_BOUNDARY_BYTES: usize = 3_570; + let large_assistant = "x".repeat(ASSISTANT_BOUNDARY_BYTES - "assistant: \n".len()); context.record_assistant_message( RetainedUserMessage { phase: None, origin: codex_history::UserInputOrigin::User, turn_id: "large".to_owned(), message_id: Some("large".to_owned()), - text: "x".repeat(3_570 - "assistant: \n".len()), + text: large_assistant.clone(), complete: true, }, RetainedInputSource::Local(Some(2)), &uuid_v7_item_id_generator_for_test(), ); - let rendered = render_retained_instructions( - &context, - retained_assistant_message, - crate::ContextTarget::Async, - ); + let rendered = render_retained_instructions(&context, retained_assistant_message); let history = RetainedTranscript { retained: context, messages: Vec::new(), @@ -249,6 +240,18 @@ fn ordinary_exchanges_keep_roles_and_drop_assistant_context_before_restrictions( .iter() .any(|fragment| fragment.content.contains("Yes, staging only.")) ); + for profile in [ + crate::ContextProfile::synchronous(), + crate::ContextProfile::asynchronous(), + ] { + let elastic = + transcript_context_with_budgeting(&history, &profile, crate::ContextBudgeting::Elastic); + assert!( + serde_json::to_string(&elastic.into_messages()) + .unwrap() + .contains(&large_assistant) + ); + } } #[test] @@ -267,6 +270,8 @@ fn legacy_verified_answers_keep_distinct_source_order() { "Retained source order: 1\nassistant: Still publish?\nuser: No.\n".to_owned(), ] ); + const ANSWER_REPETITIONS: usize = 200; + let large_answer = "Only if conditions hold. ".repeat(ANSWER_REPETITIONS); context.record( &RetainedContextEvent::VerifiedAnswer { answer: VerifiedAnswer { @@ -274,7 +279,7 @@ fn legacy_verified_answers_keep_distinct_source_order() { call_id: "large".into(), questions: vec![VerifiedQuestionAnswer { question: "Publish?".into(), - answer: "Only if conditions hold. ".repeat(200), + answer: large_answer.clone(), }], }, acceptance_order: None, @@ -284,24 +289,44 @@ fn legacy_verified_answers_keep_distinct_source_order() { let raw = crate::collect_verified_answer_candidates(&context); let old = crate::render_verified_answers(&context); assert!(!old.complete); - let [actual, expected] = - [raw, old.fragments.into_iter().map(Into::into).collect()].map(|items| { + for profile in [ + crate::ContextProfile::synchronous(), + crate::ContextProfile::asynchronous(), + ] { + let [actual, expected, elastic] = [ + (raw.clone(), crate::ContextBudgeting::Fixed), + ( + old.fragments.iter().cloned().map(Into::into).collect(), + crate::ContextBudgeting::Fixed, + ), + (raw.clone(), crate::ContextBudgeting::Elastic), + ] + .map(|(items, budgeting)| { crate::prepare_context( std::borrow::Cow::Owned(crate::CollectedContext { sections: vec![ContextSection::TrustedUserAnswers { items }], }), - &crate::ContextProfile::synchronous(), + &profile, /*entry_number_offset*/ 0, - crate::ContextPresentation::SyncFull { - session_id: "worker", + match profile.target { + crate::ContextTarget::Sync => crate::ContextPresentation::SyncFull { + session_id: "worker", + }, + crate::ContextTarget::Async => crate::ContextPresentation::Async, }, - crate::ContextBudgeting::Fixed, + budgeting, &[], ) .unwrap() .into_messages() }); - assert_eq!(actual, expected); + assert_eq!(actual, expected); + assert!( + serde_json::to_string(&elastic) + .unwrap() + .contains(&large_answer) + ); + } } #[test] @@ -472,7 +497,15 @@ fn transcript_context( history: &RetainedTranscript, profile: &crate::ContextProfile, ) -> ComposedContext { - let answers = crate::render_verified_answers(&history.retained); + transcript_context_with_budgeting(history, profile, crate::ContextBudgeting::Fixed) +} + +fn transcript_context_with_budgeting( + history: &RetainedTranscript, + profile: &crate::ContextProfile, + budgeting: crate::ContextBudgeting, +) -> ComposedContext { + let answers = crate::collect_verified_answer_candidates(&history.retained); let collected = crate::default_registry() .prepare(&SectionInput { target: profile.target, @@ -481,11 +514,7 @@ fn transcript_context( is_selected_exec_output: None, root_conversation: &[], root_conversation_candidates: None, - trusted_user_answers: &answers - .fragments - .into_iter() - .map(Into::into) - .collect::>(), + trusted_user_answers: &answers, planned_action: None, permissions: None, previous_reviews: None, @@ -506,7 +535,7 @@ fn transcript_context( profile, /*entry_number_offset*/ 0, presentation, - crate::ContextBudgeting::Fixed, + budgeting, &[], ) .unwrap()