mirror of
https://github.com/openai/codex.git
synced 2026-09-05 15:18:41 +00:00
Eagerly compact Guardian threads between reviews
This commit is contained in:
@@ -57,6 +57,10 @@ use super::prompt::build_guardian_prompt_items_with_parent_turn;
|
||||
use super::prompt::guardian_policy_prompt;
|
||||
use super::prompt::guardian_policy_prompt_with_config;
|
||||
|
||||
mod eager_compaction;
|
||||
|
||||
use eager_compaction::GuardianEagerCompaction;
|
||||
|
||||
const GUARDIAN_INTERRUPT_DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum GuardianReviewSessionOutcome {
|
||||
@@ -105,6 +109,7 @@ struct GuardianReviewSession {
|
||||
cancel_token: CancellationToken,
|
||||
reuse_key: GuardianReviewSessionReuseKey,
|
||||
review_lock: Semaphore,
|
||||
eager_compaction: GuardianEagerCompaction,
|
||||
state: Mutex<GuardianReviewState>,
|
||||
}
|
||||
|
||||
@@ -215,6 +220,7 @@ pub(crate) fn prompt_cache_key_override_for_review_session(
|
||||
impl GuardianReviewSession {
|
||||
async fn shutdown(&self) {
|
||||
self.cancel_token.cancel();
|
||||
self.wait_for_eager_compaction().await;
|
||||
let _ = self.codex.shutdown_and_wait().await;
|
||||
}
|
||||
|
||||
@@ -408,6 +414,16 @@ impl GuardianReviewSessionManager {
|
||||
.await;
|
||||
}
|
||||
|
||||
if let Err(outcome) = run_before_review_deadline(
|
||||
deadline,
|
||||
params.external_cancel.as_ref(),
|
||||
trunk.wait_for_eager_compaction(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return (outcome, GuardianReviewAnalyticsResult::without_session());
|
||||
}
|
||||
|
||||
let trunk_guard = match trunk.review_lock.try_acquire() {
|
||||
Ok(trunk_guard) => trunk_guard,
|
||||
Err(_) => {
|
||||
@@ -435,6 +451,7 @@ impl GuardianReviewSessionManager {
|
||||
.await;
|
||||
if keep_review_session && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) {
|
||||
trunk.refresh_last_committed_fork_snapshot().await;
|
||||
trunk.schedule_eager_compaction().await;
|
||||
}
|
||||
drop(trunk_guard);
|
||||
|
||||
@@ -459,6 +476,7 @@ impl GuardianReviewSessionManager {
|
||||
codex,
|
||||
cancel_token: CancellationToken::new(),
|
||||
review_lock: Semaphore::new(/*permits*/ 1),
|
||||
eager_compaction: GuardianEagerCompaction::default(),
|
||||
state: Mutex::new(GuardianReviewState {
|
||||
prior_review_count: 0,
|
||||
last_reviewed_transcript_cursor: None,
|
||||
@@ -482,6 +500,7 @@ impl GuardianReviewSessionManager {
|
||||
codex,
|
||||
cancel_token: CancellationToken::new(),
|
||||
review_lock: Semaphore::new(/*permits*/ 1),
|
||||
eager_compaction: GuardianEagerCompaction::default(),
|
||||
state: Mutex::new(GuardianReviewState {
|
||||
prior_review_count: 0,
|
||||
last_reviewed_transcript_cursor: None,
|
||||
@@ -636,6 +655,7 @@ async fn spawn_guardian_review_session(
|
||||
cancel_token,
|
||||
reuse_key,
|
||||
review_lock: Semaphore::new(/*permits*/ 1),
|
||||
eager_compaction: GuardianEagerCompaction::default(),
|
||||
state: Mutex::new(GuardianReviewState {
|
||||
prior_review_count,
|
||||
last_reviewed_transcript_cursor: initial_transcript_cursor,
|
||||
@@ -1115,6 +1135,7 @@ mod tests {
|
||||
cancel_token: CancellationToken::new(),
|
||||
reuse_key,
|
||||
review_lock: Semaphore::new(/*permits*/ 1),
|
||||
eager_compaction: GuardianEagerCompaction::default(),
|
||||
state: Mutex::new(GuardianReviewState {
|
||||
prior_review_count: 0,
|
||||
last_reviewed_transcript_cursor: None,
|
||||
|
||||
154
codex-rs/core/src/guardian/review_session/eager_compaction.rs
Normal file
154
codex-rs/core/src/guardian/review_session/eager_compaction.rs
Normal file
@@ -0,0 +1,154 @@
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::watch;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::guardian::GUARDIAN_REVIEW_TIMEOUT;
|
||||
use crate::session::turn;
|
||||
|
||||
use super::GuardianReviewSession;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "eager_compaction_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct GuardianEagerCompaction {
|
||||
completion: Mutex<Option<watch::Receiver<bool>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Eq, PartialEq)]
|
||||
enum EagerCompactionRunOutcome {
|
||||
Completed,
|
||||
Cancelled,
|
||||
TimedOut,
|
||||
}
|
||||
|
||||
struct EagerCompactionRun {
|
||||
completion: watch::Sender<bool>,
|
||||
}
|
||||
|
||||
impl EagerCompactionRun {
|
||||
async fn run_bounded<F>(
|
||||
self,
|
||||
cancel_token: &CancellationToken,
|
||||
timeout: Duration,
|
||||
maintenance: F,
|
||||
) -> EagerCompactionRunOutcome
|
||||
where
|
||||
F: Future<Output = ()>,
|
||||
{
|
||||
let outcome = tokio::select! {
|
||||
_ = cancel_token.cancelled() => EagerCompactionRunOutcome::Cancelled,
|
||||
result = tokio::time::timeout(timeout, maintenance) => {
|
||||
if result.is_ok() {
|
||||
EagerCompactionRunOutcome::Completed
|
||||
} else {
|
||||
EagerCompactionRunOutcome::TimedOut
|
||||
}
|
||||
}
|
||||
};
|
||||
self.completion.send_replace(true);
|
||||
outcome
|
||||
}
|
||||
}
|
||||
|
||||
impl GuardianEagerCompaction {
|
||||
async fn begin(&self) -> Option<EagerCompactionRun> {
|
||||
let mut completion = self.completion.lock().await;
|
||||
if let Some(receiver) = completion.as_ref()
|
||||
&& !*receiver.borrow()
|
||||
&& receiver.has_changed().is_ok()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let (sender, receiver) = watch::channel(false);
|
||||
*completion = Some(receiver);
|
||||
Some(EagerCompactionRun { completion: sender })
|
||||
}
|
||||
|
||||
async fn wait(&self) {
|
||||
let Some(mut completion) = self.completion.lock().await.clone() else {
|
||||
return;
|
||||
};
|
||||
if *completion.borrow() {
|
||||
return;
|
||||
}
|
||||
while completion.changed().await.is_ok() {
|
||||
if *completion.borrow() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GuardianReviewSession {
|
||||
pub(super) async fn schedule_eager_compaction(self: &Arc<Self>) {
|
||||
let turn_context = self.codex.session.new_default_turn().await;
|
||||
if !turn::auto_compact_needed(self.codex.session.as_ref(), turn_context.as_ref()).await {
|
||||
return;
|
||||
}
|
||||
let Some(run) = self.eager_compaction.begin().await else {
|
||||
return;
|
||||
};
|
||||
|
||||
let review_session = Arc::clone(self);
|
||||
drop(tokio::spawn(async move {
|
||||
let cancel_token = review_session.cancel_token.clone();
|
||||
let outcome = run
|
||||
.run_bounded(
|
||||
&cancel_token,
|
||||
GUARDIAN_REVIEW_TIMEOUT,
|
||||
review_session.run_eager_compaction(turn_context),
|
||||
)
|
||||
.await;
|
||||
if outcome == EagerCompactionRunOutcome::TimedOut {
|
||||
warn!(
|
||||
guardian_thread_id = %review_session.codex.session.thread_id,
|
||||
"eager guardian maintenance timed out after {GUARDIAN_REVIEW_TIMEOUT:?}"
|
||||
);
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
pub(super) async fn wait_for_eager_compaction(&self) {
|
||||
self.eager_compaction.wait().await;
|
||||
}
|
||||
|
||||
async fn run_eager_compaction(
|
||||
self: &Arc<Self>,
|
||||
turn_context: Arc<crate::session::turn_context::TurnContext>,
|
||||
) {
|
||||
let Ok(review_guard) = self.review_lock.acquire().await else {
|
||||
return;
|
||||
};
|
||||
|
||||
let mut client_session = self.codex.session.services.model_client.new_session();
|
||||
let compact_result = turn::run_pre_turn_auto_compact_if_needed(
|
||||
&self.codex.session,
|
||||
&turn_context,
|
||||
&mut client_session,
|
||||
)
|
||||
.await;
|
||||
|
||||
match compact_result {
|
||||
Ok(true) => {
|
||||
self.refresh_last_committed_fork_snapshot().await;
|
||||
}
|
||||
Ok(false) => {}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
guardian_thread_id = %self.codex.session.thread_id,
|
||||
"eager guardian compaction failed: {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
drop(review_guard);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
use std::future;
|
||||
use std::time::Duration;
|
||||
|
||||
use pretty_assertions::assert_eq;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::EagerCompactionRunOutcome;
|
||||
use super::GuardianEagerCompaction;
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancellation_releases_eager_compaction_waiters() {
|
||||
let eager_compaction = GuardianEagerCompaction::default();
|
||||
let run = eager_compaction.begin().await.expect("start maintenance");
|
||||
let cancel_token = CancellationToken::new();
|
||||
cancel_token.cancel();
|
||||
|
||||
let outcome = run
|
||||
.run_bounded(
|
||||
&cancel_token,
|
||||
Duration::from_secs(/*secs*/ 1),
|
||||
future::pending(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(outcome, EagerCompactionRunOutcome::Cancelled);
|
||||
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), eager_compaction.wait())
|
||||
.await
|
||||
.expect("waiter should be released after cancellation");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn total_timeout_releases_eager_compaction_waiters() {
|
||||
let eager_compaction = GuardianEagerCompaction::default();
|
||||
let run = eager_compaction.begin().await.expect("start maintenance");
|
||||
let cancel_token = CancellationToken::new();
|
||||
|
||||
let outcome = run
|
||||
.run_bounded(
|
||||
&cancel_token,
|
||||
Duration::from_millis(/*millis*/ 10),
|
||||
future::pending(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(outcome, EagerCompactionRunOutcome::TimedOut);
|
||||
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), eager_compaction.wait())
|
||||
.await
|
||||
.expect("waiter should be released after the total maintenance timeout");
|
||||
}
|
||||
@@ -57,6 +57,7 @@ use core_test_support::context_snapshot;
|
||||
use core_test_support::context_snapshot::ContextSnapshotOptions;
|
||||
use core_test_support::responses::ev_assistant_message;
|
||||
use core_test_support::responses::ev_completed;
|
||||
use core_test_support::responses::ev_completed_with_tokens;
|
||||
use core_test_support::responses::ev_response_created;
|
||||
use core_test_support::responses::mount_response_sequence;
|
||||
use core_test_support::responses::mount_sse_once;
|
||||
@@ -2153,6 +2154,296 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow:
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn configure_guardian_eager_compaction_test(turn: &mut Arc<TurnContext>) {
|
||||
const AUTO_COMPACT_TOKEN_LIMIT: i64 = 200_000;
|
||||
const COMPACTION_STREAM_RETRIES: u64 = 0;
|
||||
const REQUEST_RETRIES: u64 = 0;
|
||||
|
||||
let mut config = (*turn.config).clone();
|
||||
config.model_auto_compact_token_limit = Some(AUTO_COMPACT_TOKEN_LIMIT);
|
||||
config.model_provider.request_max_retries = Some(REQUEST_RETRIES);
|
||||
config.model_provider.stream_max_retries = Some(COMPACTION_STREAM_RETRIES);
|
||||
config.model_provider.supports_websockets = false;
|
||||
let _ = config.features.enable(Feature::RemoteCompactionV2);
|
||||
Arc::get_mut(turn)
|
||||
.expect("guardian test turn should be uniquely owned")
|
||||
.config = Arc::new(config);
|
||||
}
|
||||
|
||||
fn request_is_remote_compaction(request: &[u8]) -> anyhow::Result<bool> {
|
||||
let body = serde_json::from_slice::<serde_json::Value>(request)?;
|
||||
Ok(body["input"].as_array().is_some_and(|input| {
|
||||
input.iter().any(|item| {
|
||||
item.get("type").and_then(serde_json::Value::as_str) == Some("compaction_trigger")
|
||||
})
|
||||
}))
|
||||
}
|
||||
|
||||
fn eager_compaction_guardian_assessment(rationale: &str) -> String {
|
||||
serde_json::json!({
|
||||
"risk_level": "low",
|
||||
"user_authorization": "high",
|
||||
"outcome": "allow",
|
||||
"rationale": rationale,
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
async fn run_initial_eager_compaction_review(
|
||||
session: &Arc<Session>,
|
||||
turn: &Arc<TurnContext>,
|
||||
) -> (
|
||||
GuardianReviewOutcome,
|
||||
codex_analytics::GuardianReviewAnalyticsResult,
|
||||
) {
|
||||
run_guardian_review_session_for_test(
|
||||
Arc::clone(session),
|
||||
Arc::clone(turn),
|
||||
GuardianApprovalRequest::Shell {
|
||||
id: "shell-ca-540-1".to_string(),
|
||||
command: vec!["git".to_string(), "status".to_string()],
|
||||
cwd: test_path_buf("/repo/codex-rs/core").abs(),
|
||||
sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault,
|
||||
additional_permissions: None,
|
||||
justification: Some("Inspect repository state.".to_string()),
|
||||
},
|
||||
/*retry_reason*/ None,
|
||||
guardian_output_schema(),
|
||||
/*external_cancel*/ None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn run_followup_eager_compaction_review(
|
||||
session: &Arc<Session>,
|
||||
turn: &Arc<TurnContext>,
|
||||
retry_reason: &str,
|
||||
) -> (
|
||||
GuardianReviewOutcome,
|
||||
codex_analytics::GuardianReviewAnalyticsResult,
|
||||
) {
|
||||
run_guardian_review_session_for_test(
|
||||
Arc::clone(session),
|
||||
Arc::clone(turn),
|
||||
GuardianApprovalRequest::Shell {
|
||||
id: "shell-ca-540-2".to_string(),
|
||||
command: vec!["git".to_string(), "push".to_string()],
|
||||
cwd: test_path_buf("/repo/codex-rs/core").abs(),
|
||||
sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault,
|
||||
additional_permissions: None,
|
||||
justification: Some("Push the reviewed change.".to_string()),
|
||||
},
|
||||
Some(retry_reason.to_string()),
|
||||
guardian_output_schema(),
|
||||
/*external_cancel*/ None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn guardian_review_waits_until_eager_compaction_snapshot_is_committed() -> anyhow::Result<()>
|
||||
{
|
||||
const REQUEST_WAIT_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5);
|
||||
const BLOCKED_REVIEW_OBSERVATION: Duration = Duration::from_millis(/*millis*/ 100);
|
||||
const FIRST_REVIEW_TOTAL_TOKENS: i64 = 500_000;
|
||||
const COMPACTION_TOTAL_TOKENS: i64 = 50;
|
||||
|
||||
let first_assessment = eager_compaction_guardian_assessment("first guardian rationale");
|
||||
let second_assessment = eager_compaction_guardian_assessment("second guardian rationale");
|
||||
let (compaction_tx, compaction_rx) = tokio::sync::oneshot::channel();
|
||||
let (second_review_tx, second_review_rx) = tokio::sync::oneshot::channel();
|
||||
let (server, _) = start_streaming_sse_server(vec![
|
||||
vec![StreamingSseChunk {
|
||||
gate: None,
|
||||
body: sse(vec![
|
||||
ev_response_created("resp-guardian-1"),
|
||||
ev_assistant_message("msg-guardian-1", &first_assessment),
|
||||
ev_completed_with_tokens("resp-guardian-1", FIRST_REVIEW_TOTAL_TOKENS),
|
||||
]),
|
||||
}],
|
||||
vec![StreamingSseChunk {
|
||||
gate: Some(compaction_rx),
|
||||
body: sse(vec![
|
||||
ev_response_created("resp-compact-success"),
|
||||
serde_json::json!({
|
||||
"type": "response.output_item.done",
|
||||
"item": {
|
||||
"type": "compaction",
|
||||
"encrypted_content": "CA540_COMPACTED_CONTEXT",
|
||||
}
|
||||
}),
|
||||
ev_completed_with_tokens("resp-compact-success", COMPACTION_TOTAL_TOKENS),
|
||||
]),
|
||||
}],
|
||||
vec![StreamingSseChunk {
|
||||
gate: Some(second_review_rx),
|
||||
body: sse(vec![
|
||||
ev_response_created("resp-guardian-2"),
|
||||
ev_assistant_message("msg-guardian-2", &second_assessment),
|
||||
ev_completed("resp-guardian-2"),
|
||||
]),
|
||||
}],
|
||||
])
|
||||
.await;
|
||||
|
||||
let (session, mut turn) = guardian_test_session_and_turn_with_base_url(server.uri()).await;
|
||||
configure_guardian_eager_compaction_test(&mut turn);
|
||||
seed_guardian_parent_history(&session, &turn).await;
|
||||
|
||||
let first_outcome = tokio::time::timeout(
|
||||
REQUEST_WAIT_TIMEOUT,
|
||||
run_initial_eager_compaction_review(&session, &turn),
|
||||
)
|
||||
.await?;
|
||||
let (GuardianReviewOutcome::Completed(_), _) = first_outcome else {
|
||||
panic!("expected first guardian assessment");
|
||||
};
|
||||
tokio::time::timeout(REQUEST_WAIT_TIMEOUT, server.wait_for_request_count(2)).await?;
|
||||
|
||||
let session_for_second = Arc::clone(&session);
|
||||
let turn_for_second = Arc::clone(&turn);
|
||||
let second_review = tokio::spawn(async move {
|
||||
run_followup_eager_compaction_review(
|
||||
&session_for_second,
|
||||
&turn_for_second,
|
||||
"Follow-up approval on the reused guardian thread.",
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
assert!(
|
||||
tokio::time::timeout(BLOCKED_REVIEW_OBSERVATION, server.wait_for_request_count(3))
|
||||
.await
|
||||
.is_err(),
|
||||
"the next guardian request must not start while eager compaction is in flight"
|
||||
);
|
||||
assert!(!second_review.is_finished());
|
||||
|
||||
compaction_tx
|
||||
.send(())
|
||||
.expect("eager compaction response gate should still be open");
|
||||
tokio::time::timeout(REQUEST_WAIT_TIMEOUT, server.wait_for_request_count(3)).await?;
|
||||
|
||||
let committed_rollout_items = session
|
||||
.guardian_review_session
|
||||
.committed_fork_rollout_items_for_test()
|
||||
.await
|
||||
.expect("committed guardian fork snapshot");
|
||||
assert!(
|
||||
committed_rollout_items.iter().any(|item| matches!(
|
||||
item,
|
||||
RolloutItem::Compacted(compacted) if compacted.replacement_history.is_some()
|
||||
)),
|
||||
"eager compaction must refresh the fork snapshot before releasing the next review"
|
||||
);
|
||||
assert!(!second_review.is_finished());
|
||||
|
||||
second_review_tx
|
||||
.send(())
|
||||
.expect("second guardian review gate should still be open");
|
||||
let second_outcome = tokio::time::timeout(REQUEST_WAIT_TIMEOUT, second_review).await??;
|
||||
let (GuardianReviewOutcome::Completed(_), second_metadata) = second_outcome else {
|
||||
panic!("expected second guardian assessment");
|
||||
};
|
||||
assert!(matches!(
|
||||
second_metadata.guardian_session_kind,
|
||||
Some(codex_analytics::GuardianReviewSessionKind::TrunkReused)
|
||||
));
|
||||
|
||||
let requests = server.requests().await;
|
||||
assert_eq!(requests.len(), 3);
|
||||
assert!(request_is_remote_compaction(&requests[1])?);
|
||||
assert!(!request_is_remote_compaction(&requests[2])?);
|
||||
|
||||
server.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn guardian_eager_compaction_failure_falls_back_to_pre_turn_compaction() -> anyhow::Result<()>
|
||||
{
|
||||
const REQUEST_WAIT_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5);
|
||||
const FIRST_REVIEW_TOTAL_TOKENS: i64 = 500_000;
|
||||
const COMPACTION_TOTAL_TOKENS: i64 = 50;
|
||||
|
||||
let first_assessment = eager_compaction_guardian_assessment("first guardian rationale");
|
||||
let second_assessment = eager_compaction_guardian_assessment("second guardian rationale");
|
||||
let (server, _) = start_streaming_sse_server(vec![
|
||||
vec![StreamingSseChunk {
|
||||
gate: None,
|
||||
body: sse(vec![
|
||||
ev_response_created("resp-guardian-1"),
|
||||
ev_assistant_message("msg-guardian-1", &first_assessment),
|
||||
ev_completed_with_tokens("resp-guardian-1", FIRST_REVIEW_TOTAL_TOKENS),
|
||||
]),
|
||||
}],
|
||||
vec![StreamingSseChunk {
|
||||
gate: None,
|
||||
body: sse(vec![ev_response_created("resp-eager-compact-failed")]),
|
||||
}],
|
||||
vec![StreamingSseChunk {
|
||||
gate: None,
|
||||
body: sse(vec![
|
||||
ev_response_created("resp-pre-turn-compact-success"),
|
||||
serde_json::json!({
|
||||
"type": "response.output_item.done",
|
||||
"item": {
|
||||
"type": "compaction",
|
||||
"encrypted_content": "CA540_RETRY_COMPACTED_CONTEXT",
|
||||
}
|
||||
}),
|
||||
ev_completed_with_tokens("resp-pre-turn-compact-success", COMPACTION_TOTAL_TOKENS),
|
||||
]),
|
||||
}],
|
||||
vec![StreamingSseChunk {
|
||||
gate: None,
|
||||
body: sse(vec![
|
||||
ev_response_created("resp-guardian-2"),
|
||||
ev_assistant_message("msg-guardian-2", &second_assessment),
|
||||
ev_completed("resp-guardian-2"),
|
||||
]),
|
||||
}],
|
||||
])
|
||||
.await;
|
||||
|
||||
let (session, mut turn) = guardian_test_session_and_turn_with_base_url(server.uri()).await;
|
||||
configure_guardian_eager_compaction_test(&mut turn);
|
||||
seed_guardian_parent_history(&session, &turn).await;
|
||||
|
||||
let first_outcome = tokio::time::timeout(
|
||||
REQUEST_WAIT_TIMEOUT,
|
||||
run_initial_eager_compaction_review(&session, &turn),
|
||||
)
|
||||
.await?;
|
||||
let (GuardianReviewOutcome::Completed(_), _) = first_outcome else {
|
||||
panic!("expected first guardian assessment");
|
||||
};
|
||||
tokio::time::timeout(REQUEST_WAIT_TIMEOUT, server.wait_for_request_count(2)).await?;
|
||||
|
||||
let second_outcome = run_followup_eager_compaction_review(
|
||||
&session,
|
||||
&turn,
|
||||
"Retry after eager compaction failure.",
|
||||
)
|
||||
.await;
|
||||
let (GuardianReviewOutcome::Completed(_), second_metadata) = second_outcome else {
|
||||
panic!("expected second guardian assessment");
|
||||
};
|
||||
assert!(matches!(
|
||||
second_metadata.guardian_session_kind,
|
||||
Some(codex_analytics::GuardianReviewSessionKind::TrunkReused)
|
||||
));
|
||||
|
||||
let requests = server.requests().await;
|
||||
assert_eq!(requests.len(), 4);
|
||||
assert!(request_is_remote_compaction(&requests[1])?);
|
||||
assert!(request_is_remote_compaction(&requests[2])?);
|
||||
assert!(!request_is_remote_compaction(&requests[3])?);
|
||||
|
||||
server.shutdown().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow::Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
@@ -802,22 +802,37 @@ async fn run_pre_sampling_compact(
|
||||
client_session: &mut ModelClientSession,
|
||||
) -> CodexResult<()> {
|
||||
maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?;
|
||||
let token_status = auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await;
|
||||
// Compact if the configured auto-compaction budget or usable context window is exhausted.
|
||||
if token_status.token_limit_reached {
|
||||
run_auto_compact(
|
||||
sess,
|
||||
turn_context,
|
||||
client_session,
|
||||
InitialContextInjection::DoNotInject,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionPhase::PreTurn,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
run_pre_turn_auto_compact_if_needed(sess, turn_context, client_session).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn auto_compact_needed(sess: &Session, turn_context: &TurnContext) -> bool {
|
||||
auto_compact_token_status(sess, turn_context)
|
||||
.await
|
||||
.token_limit_reached
|
||||
}
|
||||
|
||||
pub(crate) async fn run_pre_turn_auto_compact_if_needed(
|
||||
sess: &Arc<Session>,
|
||||
turn_context: &Arc<TurnContext>,
|
||||
client_session: &mut ModelClientSession,
|
||||
) -> CodexResult<bool> {
|
||||
if !auto_compact_needed(sess.as_ref(), turn_context.as_ref()).await {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
run_auto_compact(
|
||||
sess,
|
||||
turn_context,
|
||||
client_session,
|
||||
InitialContextInjection::DoNotInject,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionPhase::PreTurn,
|
||||
)
|
||||
.await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// Returns true only when both turns declare compaction compatibility hashes and they differ.
|
||||
/// A missing hash does not provide enough information to trigger compaction.
|
||||
fn comp_hash_changed(previous: Option<&str>, current: Option<&str>) -> bool {
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
use anyhow::Result;
|
||||
use codex_core::config::Constrained;
|
||||
use codex_core::sandboxing::SandboxPermissions;
|
||||
use codex_features::Feature;
|
||||
use codex_protocol::config_types::ApprovalsReviewer;
|
||||
use codex_protocol::protocol::AskForApproval;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
@@ -12,6 +13,7 @@ use codex_protocol::user_input::UserInput;
|
||||
use core_test_support::fs_wait;
|
||||
use core_test_support::responses::ev_assistant_message;
|
||||
use core_test_support::responses::ev_completed;
|
||||
use core_test_support::responses::ev_completed_with_tokens;
|
||||
use core_test_support::responses::ev_function_call;
|
||||
use core_test_support::responses::ev_response_created;
|
||||
use core_test_support::responses::mount_sse_sequence;
|
||||
@@ -30,6 +32,171 @@ use std::os::unix::fs::PermissionsExt;
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn guardian_compacts_between_reviews_before_the_next_request() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
skip_if_sandbox!(Ok(()));
|
||||
|
||||
const FIRST_REVIEW_TOTAL_TOKENS: i64 = 500_000;
|
||||
const COMPACTION_TOTAL_TOKENS: i64 = 50;
|
||||
const AUTO_COMPACT_TOKEN_LIMIT: i64 = 200_000;
|
||||
const REQUEST_RETRIES: u64 = 0;
|
||||
const STREAM_RETRIES: u64 = 0;
|
||||
|
||||
let server = start_mock_server().await;
|
||||
let approval_policy = AskForApproval::OnRequest;
|
||||
let sandbox_policy = SandboxPolicy::WorkspaceWrite {
|
||||
writable_roots: vec![],
|
||||
network_access: false,
|
||||
exclude_tmpdir_env_var: true,
|
||||
exclude_slash_tmp: true,
|
||||
};
|
||||
let sandbox_policy_for_config = sandbox_policy.clone();
|
||||
let mut builder = test_codex().with_config(move |config| {
|
||||
config.permissions.approval_policy = Constrained::allow_any(approval_policy);
|
||||
config
|
||||
.set_legacy_sandbox_policy(sandbox_policy_for_config)
|
||||
.expect("set sandbox policy");
|
||||
config.model_auto_compact_token_limit = Some(AUTO_COMPACT_TOKEN_LIMIT);
|
||||
config.model_provider.request_max_retries = Some(REQUEST_RETRIES);
|
||||
config.model_provider.stream_max_retries = Some(STREAM_RETRIES);
|
||||
config.model_provider.supports_websockets = false;
|
||||
config
|
||||
.features
|
||||
.enable(Feature::RemoteCompactionV2)
|
||||
.expect("enable remote compaction v2");
|
||||
});
|
||||
let test = builder.build(&server).await?;
|
||||
|
||||
let first_justification = "Approve the first eager-compaction command.";
|
||||
let second_justification = "Approve the second eager-compaction command.";
|
||||
let first_args = json!({
|
||||
"cmd": "sleep 0.2; printf first",
|
||||
"yield_time_ms": 1_000_u64,
|
||||
"sandbox_permissions": SandboxPermissions::RequireEscalated,
|
||||
"justification": first_justification,
|
||||
});
|
||||
let second_args = json!({
|
||||
"cmd": "printf second",
|
||||
"yield_time_ms": 1_000_u64,
|
||||
"sandbox_permissions": SandboxPermissions::RequireEscalated,
|
||||
"justification": second_justification,
|
||||
});
|
||||
let guardian_assessment = json!({
|
||||
"risk_level": "low",
|
||||
"user_authorization": "high",
|
||||
"outcome": "allow",
|
||||
"rationale": "The command writes a bounded marker file requested by the user.",
|
||||
})
|
||||
.to_string();
|
||||
let responses = mount_sse_sequence(
|
||||
&server,
|
||||
vec![
|
||||
sse(vec![
|
||||
ev_response_created("resp-parent-first-tool"),
|
||||
ev_function_call(
|
||||
"exec-first",
|
||||
"exec_command",
|
||||
&serde_json::to_string(&first_args)?,
|
||||
),
|
||||
ev_completed("resp-parent-first-tool"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-guardian-first"),
|
||||
ev_assistant_message("msg-guardian-first", &guardian_assessment),
|
||||
ev_completed_with_tokens("resp-guardian-first", FIRST_REVIEW_TOTAL_TOKENS),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-guardian-compact"),
|
||||
json!({
|
||||
"type": "response.output_item.done",
|
||||
"item": {
|
||||
"type": "compaction",
|
||||
"encrypted_content": "EAGER_GUARDIAN_COMPACTED_CONTEXT",
|
||||
}
|
||||
}),
|
||||
ev_completed_with_tokens("resp-guardian-compact", COMPACTION_TOTAL_TOKENS),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-parent-second-tool"),
|
||||
ev_function_call(
|
||||
"exec-second",
|
||||
"exec_command",
|
||||
&serde_json::to_string(&second_args)?,
|
||||
),
|
||||
ev_completed("resp-parent-second-tool"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-guardian-second"),
|
||||
ev_assistant_message("msg-guardian-second", &guardian_assessment),
|
||||
ev_completed("resp-guardian-second"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-parent-done"),
|
||||
ev_assistant_message("msg-parent-done", "done"),
|
||||
ev_completed("resp-parent-done"),
|
||||
]),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
test.codex
|
||||
.submit(Op::UserInput {
|
||||
items: vec![UserInput::Text {
|
||||
text: "run both commands that require Guardian review".into(),
|
||||
text_elements: Vec::new(),
|
||||
}],
|
||||
final_output_json_schema: None,
|
||||
responsesapi_client_metadata: None,
|
||||
additional_context: Default::default(),
|
||||
thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
|
||||
environments: Some(local_selections(test.config.cwd.clone())),
|
||||
approval_policy: Some(approval_policy),
|
||||
approvals_reviewer: Some(ApprovalsReviewer::AutoReview),
|
||||
sandbox_policy: Some(sandbox_policy),
|
||||
..Default::default()
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
|
||||
let requests = responses.requests();
|
||||
let first_guardian_index = requests
|
||||
.iter()
|
||||
.position(|request| request.body_contains_text(first_justification))
|
||||
.expect("first Guardian request");
|
||||
let second_guardian_index = requests
|
||||
.iter()
|
||||
.position(|request| request.body_contains_text(second_justification))
|
||||
.expect("second Guardian request");
|
||||
let compaction_indexes = requests
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, request)| {
|
||||
request.body_json()["input"]
|
||||
.as_array()
|
||||
.is_some_and(|input| {
|
||||
input.iter().any(|item| {
|
||||
item.get("type").and_then(Value::as_str) == Some("compaction_trigger")
|
||||
})
|
||||
})
|
||||
.then_some(index)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(compaction_indexes.len(), 1);
|
||||
assert!(
|
||||
first_guardian_index < compaction_indexes[0]
|
||||
&& compaction_indexes[0] < second_guardian_index,
|
||||
"expected eager compaction between Guardian reviews, requests: {requests:#?}"
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn guardian_review_session_does_not_inherit_legacy_notify() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
Reference in New Issue
Block a user