diff --git a/codex-rs/ext/guardian-v3/Cargo.toml b/codex-rs/ext/guardian-v3/Cargo.toml index 7ac03126cb..1f7d899530 100644 --- a/codex-rs/ext/guardian-v3/Cargo.toml +++ b/codex-rs/ext/guardian-v3/Cargo.toml @@ -24,7 +24,7 @@ codex-thread-store = { workspace = true } codex-tools = { workspace = true } codex-utils-string = { workspace = true } serde = { workspace = true } -serde_json = { workspace = true } +serde_json = { workspace = true, features = ["raw_value"] } tokio = { workspace = true, features = ["time"] } tokio-util = { workspace = true } diff --git a/codex-rs/ext/guardian-v3/src/delivery.rs b/codex-rs/ext/guardian-v3/src/delivery.rs index da094344e3..1df87cc191 100644 --- a/codex-rs/ext/guardian-v3/src/delivery.rs +++ b/codex-rs/ext/guardian-v3/src/delivery.rs @@ -1,15 +1,24 @@ -//! Plans delivery of routed originals in source order by reading IDs from a fixed board view. +//! Delivers housekeeper posts from the reviewer's selected channels in source order. +//! A plan fetches bodies one board page at a time; only IDs and positions survive the plan. use crate::EvidenceId; use crate::local_board_options; +use crate::sampling::context_input; use anyhow::ensure; use codex_message_board::LocalBoardReadView; use codex_message_board::PageRequest; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::HistoryPosition; +use codex_protocol::protocol::TruncationPolicy; +use serde::Deserialize; +use serde::Serialize; +use serde_json::value::RawValue; use std::collections::BTreeMap; use std::collections::BTreeSet; +use std::collections::VecDeque; +use std::future::Future; /// In-memory progress for one retained reviewer conversation. A fresh conversation starts with fresh state. #[derive(Clone, Debug, PartialEq, Eq)] @@ -21,6 +30,31 @@ pub struct DeliveryState { channel_positions: BTreeMap, } +/// Exclusive to one selection and captured view. Dropping a partial plan preserves +/// successful admissions. A subsequent plan resumes from the last completed channel ranges. +pub struct DeliveryPlan<'a> { + state: &'a mut DeliveryState, + board: &'a LocalBoardReadView, + /// Source ordinal -> one immutable housekeeper post. + pending: BTreeMap, + /// Texts for the leading pending ordinals, including any that did not fit in the last batch. + prefetched: VecDeque, + channel_positions: BTreeMap, +} + +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +struct RoutedOriginal<'a> { + id: EvidenceId, + #[serde(borrow)] + original: &'a RawValue, +} + +#[derive(Serialize)] +struct DeliveryInput<'a> { + sources: Vec>, +} + impl DeliveryState { pub fn new(reviewer: ThreadId, board: SessionId) -> Self { Self { @@ -31,13 +65,13 @@ impl DeliveryState { } } - /// Maps each selected source ordinal to one board insertion sequence within the two cutoffs. - pub async fn plan( - &self, - board: &LocalBoardReadView, + /// Collects the positions of undelivered originals within the board and source cutoffs. + pub async fn plan<'a>( + &'a mut self, + board: &'a LocalBoardReadView, selected: &BTreeSet, cutoff: HistoryPosition, - ) -> anyhow::Result> { + ) -> anyhow::Result> { ensure!( self.board == board.board().identity() && self.board == SessionId::from(cutoff.thread_id), @@ -54,6 +88,7 @@ impl DeliveryState { "review cutoff predates evidence already in the conversation" ); let mut pending = BTreeMap::new(); + let mut channel_positions = BTreeMap::new(); let page_size = PageRequest::default().limit.get() as usize; for channel in selected { ensure!( @@ -63,11 +98,12 @@ impl DeliveryState { .any(|entry| &entry.channel_name == channel), "selected channel is outside the captured board view" ); - let mut position = self + let previous = self .channel_positions .get(channel) .copied() .unwrap_or_default(); + let mut position = previous; loop { let page = board .appended_request_ids(self.reviewer, channel, position) @@ -88,7 +124,82 @@ impl DeliveryState { break; } } + if position != previous { + channel_positions.insert(channel.clone(), position); + } } - Ok(pending) + Ok(DeliveryPlan { + state: self, + board, + pending, + prefetched: VecDeque::new(), + channel_positions, + }) + } +} + +impl DeliveryPlan<'_> { + /// Admits one batch of selected board posts. The callback must record the items + /// in the retained conversation before reporting success. Returns zero when done. + pub async fn deliver( + &mut self, + policy: TruncationPolicy, + admit: F, + ) -> anyhow::Result + where + F: FnOnce(Vec) -> Fut, + Fut: Future>, + { + let mut input = DeliveryInput { + sources: Vec::new(), + }; + let mut bytes = serde_json::to_string(&input)?.len(); + let page_size = PageRequest::default().limit.get() as usize; + let mut unread = self.pending.values().skip(self.prefetched.len()); + let mut count = 0; + while count < self.pending.len() { + if count == self.prefetched.len() { + let sequences: Vec<_> = unread.by_ref().take(page_size).copied().collect(); + self.prefetched.extend( + self.board + .read_post_texts(self.state.reviewer, &sequences) + .await?, + ); + } + let additional = self.prefetched[count].len() + usize::from(count != 0); + if additional > policy.byte_budget().saturating_sub(bytes) { + ensure!( + count != 0, + "original evidence exceeds the model's tool-output budget" + ); + break; + } + bytes += additional; + count += 1; + } + for (ordinal, text) in self.pending.keys().zip(&self.prefetched).take(count) { + let source: RoutedOriginal<'_> = serde_json::from_str(text)?; + ensure!( + source.id.ordinal == *ordinal, + "selected board post has a different source ordinal" + ); + input.sources.push(source); + } + if count != 0 { + let item = context_input("guardian_original_evidence", &input, policy)?; + admit(vec![item]).await?; + self.prefetched.drain(..count); + for _ in 0..count { + if let Some((ordinal, _)) = self.pending.pop_first() { + self.state.delivered.insert(ordinal); + } + } + } + if self.pending.is_empty() { + self.state + .channel_positions + .append(&mut self.channel_positions); + } + Ok(count) } } diff --git a/codex-rs/ext/guardian-v3/src/lib.rs b/codex-rs/ext/guardian-v3/src/lib.rs index 7efb9dc386..fdd6d52bcd 100644 --- a/codex-rs/ext/guardian-v3/src/lib.rs +++ b/codex-rs/ext/guardian-v3/src/lib.rs @@ -1,5 +1,5 @@ //! Guardian's durable evidence storage, owned by the main thread. -//! Reviewer tools search a fixed board view and plan delivery of the originals routed by the housekeeper. +//! Reviewer tools search a fixed board view and deliver the originals routed by the housekeeper. //! The host owns review scheduling and the lifetime of each view. //! Hosts invoke permanent-deletion cleanup only when Guardian V3 is enabled. @@ -15,6 +15,7 @@ mod retrieval_tools; mod sampling; mod selection; mod subscription_tools; +pub use delivery::DeliveryPlan; pub use delivery::DeliveryState; pub use evidence::EvidenceCursor; pub use evidence::EvidenceId; diff --git a/codex-rs/ext/guardian-v3/tests/retrieval.rs b/codex-rs/ext/guardian-v3/tests/retrieval.rs index a9349257bf..6c46d94ce6 100644 --- a/codex-rs/ext/guardian-v3/tests/retrieval.rs +++ b/codex-rs/ext/guardian-v3/tests/retrieval.rs @@ -25,6 +25,8 @@ use std::sync::atomic::AtomicUsize; #[path = "../../../message-board/tests/common/mod.rs"] mod board_host; +#[path = "retrieval/delivery.rs"] +mod delivery; #[path = "retrieval/evidence_board.rs"] mod evidence_board; #[path = "retrieval/planning.rs"] diff --git a/codex-rs/ext/guardian-v3/tests/retrieval/delivery.rs b/codex-rs/ext/guardian-v3/tests/retrieval/delivery.rs new file mode 100644 index 0000000000..e0a06f6a27 --- /dev/null +++ b/codex-rs/ext/guardian-v3/tests/retrieval/delivery.rs @@ -0,0 +1,245 @@ +//! Whole originals enter a real retained conversation without scheduling inference. + +use codex_analytics::GuardianReviewSessionKind; +use codex_core::guardian_review::prepare_review_prewarm; +use codex_extension_api::SessionIsolation; +use codex_extension_api::ToolPolicy; +use codex_guardian_v3::DeliveryState; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::TruncationPolicyConfig; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::HistoryPosition; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TruncationPolicy; +use codex_protocol::turn_input::TurnInputRequest; +use codex_protocol::user_input::UserInput; +use core_test_support::context_snapshot; +use core_test_support::context_snapshot::ContextSnapshotOptions; +use core_test_support::responses::*; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::collections::BTreeSet; +use std::sync::atomic::Ordering; + +use crate::evidence_board::EvidenceBoard; +use crate::evidence_board::routed_original; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn delivers_ordered_originals_once_after_admission_and_keeps_older_channel_evidence() +-> anyhow::Result<()> { + core_test_support::skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex() + .with_model_info_override("test-gpt-5.1-codex", |model| { + model.auto_review_model_override = Some(model.slug.clone()); + model.truncation_policy = TruncationPolicyConfig::bytes(/*limit*/ 64 * 1024); + }) + .with_config(|config| { + config.approvals_reviewer = ApprovalsReviewer::User; + config.include_environment_context = false; + }) + .build_with_auto_env(&server) + .await?; + let owner = test.session_configured.thread_id; + let prepared = prepare_review_prewarm(&test.codex).await?; + let (mut options, _state) = prepared + .thread_options(GuardianReviewSessionKind::TrunkNew, /*snapshot*/ None) + .await; + options.config.base_instructions = Some("Inspect the supplied original evidence.".into()); + options.session_source = Some(SessionSource::Internal(InternalSessionSource::Guardian)); + options + .thread_extension_init + .insert(SessionIsolation::Isolated); + options.thread_extension_init.insert(ToolPolicy { + allowed_tools: Some(Vec::new()), + ..Default::default() + }); + let spawned = test.thread_manager.start_thread(options).await?; + let reviewer = spawned.thread_id; + let board = EvidenceBoard::new(owner, reviewer, &["a", "b", "older", "large"]).await?; + let large = "é".repeat(20_001); + // The housekeeper can route one source to several channels and post out of source order. + for (ordinal, channel, output) in [ + (3, "a", "two"), + (2, "b", "one"), + (2, "a", "one"), + (1, "older", "older"), + (4, "large", large.as_str()), + ] { + board.append(ordinal, channel, output).await?; + } + let evidence = board.writer.read_view(reviewer).await?; + let cutoff = HistoryPosition { + thread_id: owner, + end_ordinal_exclusive: 6, + end_byte_offset: 64 * 1024, + }; + let mut state = DeliveryState::new(reviewer, owner.into()); + let before_admission = state.clone(); + let selected = BTreeSet::from(["a".into(), "b".into()]); + let mut plan = state.plan(&evidence, &selected, cutoff).await?; + let small = TruncationPolicy::Bytes(160); + assert!( + plan.deliver(small, |_| async { anyhow::bail!("admission unavailable") }) + .await + .is_err() + ); + drop(plan); + assert_eq!(state, before_admission); + let mut plan = state.plan(&evidence, &selected, cutoff).await?; + let response = mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("review", "Evidence received."), + ev_completed("review"), + ]), + ) + .await; + let reads = board.host.agent_path_calls.load(Ordering::SeqCst); + for _ in 0..2 { + assert_eq!( + plan.deliver(small, |items| async { + Ok(spawned.thread.inject_response_items(items).await?) + }) + .await?, + 1 + ); + } + // Two admissions share one page of board reads. + assert_eq!( + board.host.agent_path_calls.load(Ordering::SeqCst) - reads, + 1 + ); + assert!(response.requests().is_empty()); + spawned + .thread + .start_turn_if_idle(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Assess.".into(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&spawned.thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + insta::assert_snapshot!(context_snapshot::format_request_history_snapshot( + "Selected originals arrive once in source order before a later reviewer turn.", &response.requests(), + &ContextSnapshotOptions::default().rewrite_known_segments(), + ), @r###" + Scenario: Selected originals arrive once in source order before a later reviewer turn. + + ## Window 1 + -- request 1 (turn) -- + 00:message/developer: + + 01:message/developer: + + 02:function_call_output/guardian.guardian_original_evidence:{"sources":[{"id":{"ordinal":2},"original":{"type":"response_item","payload":{"type":"function_call_output","output":"one"}}}]} + 03:function_call_output/guardian.guardian_original_evidence:{"sources":[{"id":{"ordinal":3},"original":{"type":"response_item","payload":{"type":"function_call_output","output":"two"}}}]} + 04:message/user: + Assess. + "###); + // A later subscription can add an older source without resetting delivered progress. + let mut older = state + .plan(&evidence, &BTreeSet::from(["older".into()]), cutoff) + .await?; + assert_eq!( + older + .deliver(small, |items| async { + Ok(spawned.thread.inject_response_items(items).await?) + }) + .await?, + 1 + ); + let checkpoint = state.clone(); + // A whole source larger than the admission budget can be retried with more room. + let mut long = state + .plan(&evidence, &BTreeSet::from(["large".into()]), cutoff) + .await?; + assert!( + long.deliver(small, |_| async { + panic!("oversized originals must not reach admission") + }) + .await + .is_err() + ); + drop(long); + assert_eq!(state, checkpoint); + let mut long = state + .plan(&evidence, &BTreeSet::from(["large".into()]), cutoff) + .await?; + assert_eq!( + long.deliver(TruncationPolicy::Bytes(64 * 1024), |items| async { + Ok(spawned.thread.inject_response_items(items).await?) + }) + .await?, + 1 + ); + // A backdated post and a newly routed copy both arrive after the board capture. + board.host.clock.store(-1000, Ordering::SeqCst); + for (ordinal, channel, output) in [(5, "a", "late"), (3, "b", "two")] { + board.append(ordinal, channel, output).await?; + } + let mut repeated = state.plan(&evidence, &selected, cutoff).await?; + assert_eq!( + repeated + .deliver(small, |_| async { + panic!("the old view has no additional evidence") + }) + .await?, + 0 + ); + let changed = board.writer.read_view(reviewer).await?; + let mut fresh = state.plan(&changed, &selected, cutoff).await?; + assert_eq!( + fresh + .deliver(small, |items| async { + Ok(spawned.thread.inject_response_items(items).await?) + }) + .await?, + 1 + ); + assert_eq!( + fresh + .deliver(small, |_| async { + panic!("the duplicate source must not be admitted again") + }) + .await?, + 0 + ); + + let history = spawned.thread.conversation_history_snapshot().await; + let mut delivered = Vec::new(); + for item in history.items() { + if let ResponseItem::FunctionCallOutput { + name: Some(name), + output, + .. + } = item + && name == "guardian_original_evidence" + { + let text = output + .text_content() + .ok_or_else(|| anyhow::anyhow!("expected evidence text"))?; + delivered.push(serde_json::from_str::(text)?); + } + } + let expected = [ + (2, "one"), + (3, "two"), + (1, "older"), + (4, large.as_str()), + (5, "late"), + ] + .map(|(ordinal, output)| json!({"sources":[routed_original(ordinal, output)]})); + assert_eq!(delivered.as_slice(), expected.as_slice()); + assert_eq!(response.requests().len(), 1); + + spawned.thread.shutdown_and_wait().await?; + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/ext/guardian-v3/tests/retrieval/planning.rs b/codex-rs/ext/guardian-v3/tests/retrieval/planning.rs index 332164165c..5adcd6e02e 100644 --- a/codex-rs/ext/guardian-v3/tests/retrieval/planning.rs +++ b/codex-rs/ext/guardian-v3/tests/retrieval/planning.rs @@ -1,11 +1,18 @@ //! Planning reads only the selected channels and respects both the board and history cutoffs. use crate::evidence_board::EvidenceBoard; +use crate::evidence_board::routed_original; +use anyhow::Context; +use codex_guardian_v3::DeliveryPlan; use codex_guardian_v3::DeliveryState; use codex_message_board::PageRequest; use codex_protocol::ThreadId; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::HistoryPosition; +use codex_protocol::protocol::TruncationPolicy; use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; use std::collections::BTreeSet; use std::sync::atomic::Ordering; @@ -29,15 +36,21 @@ async fn plans_unique_originals_in_source_order_from_a_fixed_view() -> anyhow::R end_byte_offset: 64 * 1024, }; let selected = BTreeSet::from(["a".into(), "b".into()]); - let state = DeliveryState::new(reviewer, owner.into()); + let mut state = DeliveryState::new(reviewer, owner.into()); let reads = board.host.agent_path_calls.load(Ordering::SeqCst); - let plan = state.plan(&view, &selected, cutoff).await?; + let mut plan = state.plan(&view, &selected, cutoff).await?; // One short ID page per selected channel completes the scan. assert_eq!( board.host.agent_path_calls.load(Ordering::SeqCst) - reads, 2 ); - assert_eq!(plan.into_keys().collect::>(), [1, 3]); + assert_eq!( + admit_batch(&mut plan).await?, + json!({"sources": [ + routed_original(/*ordinal*/ 1, "one"), + routed_original(/*ordinal*/ 3, "three"), + ]}), + ); let later = board.writer.read_view(reviewer).await?; assert!(state.plan(&later, &selected, cutoff).await.is_err()); Ok(()) @@ -59,11 +72,30 @@ async fn plans_across_storage_pages_in_reverse_insertion_order() -> anyhow::Resu end_ordinal_exclusive: count, end_byte_offset: 64 * 1024, }; - let state = DeliveryState::new(reviewer, owner.into()); - let plan = state.plan(&view, &selected, cutoff).await?; + let mut state = DeliveryState::new(reviewer, owner.into()); + let mut plan = state.plan(&view, &selected, cutoff).await?; + let reads = board.host.agent_path_calls.load(Ordering::SeqCst); + let expected: Vec<_> = (0..count) + .map(|ordinal| routed_original(ordinal, "next review")) + .collect(); + assert_eq!(admit_batch(&mut plan).await?, json!({"sources": expected})); assert_eq!( - plan.into_keys().collect::>(), - (0..count).collect::>() + board.host.agent_path_calls.load(Ordering::SeqCst) - reads, + 2 ); Ok(()) } + +async fn admit_batch(plan: &mut DeliveryPlan<'_>) -> anyhow::Result { + let mut batch = Vec::new(); + plan.deliver(TruncationPolicy::Bytes(64 * 1024), |items| { + batch = items; + std::future::ready(Ok(())) + }) + .await?; + let [ResponseItem::FunctionCallOutput { output, .. }] = batch.as_slice() else { + anyhow::bail!("expected a single evidence batch"); + }; + let text = output.text_content().context("expected evidence text")?; + Ok(serde_json::from_str(text)?) +} diff --git a/codex-rs/message-board/src/local/catalog.rs b/codex-rs/message-board/src/local/catalog.rs index b4f9db8799..8d54234c26 100644 --- a/codex-rs/message-board/src/local/catalog.rs +++ b/codex-rs/message-board/src/local/catalog.rs @@ -87,6 +87,35 @@ impl LocalBoardReadView { .await .map_err(storage_error) } + + /// Reads one page of whole posts in the requested sequence order. Every post must + /// belong to this board and fall within the captured boundary. + pub async fn read_post_texts( + &self, + caller: ThreadId, + sequences: &[i64], + ) -> Result> { + self.board.host.agent_path(caller).await?; + if sequences.len() > crate::PageRequest::default().limit.get() as usize { + return Err(invalid("requested posts exceed one board page")); + } + // Drive primary-key lookups from the requested page, not from all posts on the board. + let texts = sqlx::query_scalar::<_, String>( + "SELECT json_extract(p.payload, '$.text') FROM json_each(?) AS requested + CROSS JOIN posts AS p + WHERE p.seq=requested.value AND p.board=? AND p.seq<=? ORDER BY requested.key", + ) + .bind(serde_json::to_string(sequences).map_err(storage_error)?) + .bind(self.board.identity.to_string()) + .bind(self.view.post_cutoff) + .fetch_all(self.board.pool.as_ref()) + .await + .map_err(storage_error)?; + if texts.len() != sequences.len() { + return Err(invalid("post not found in this board view")); + } + Ok(texts) + } } pub(super) struct ReadView {