mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
Wire background compaction into the turn lifecycle (#53800)
## What changed - Launch background compaction at eligible pre-turn, mid-turn, and post-turn checkpoints when `AsyncCompaction` is enabled, capturing the current world state and client routing state. - Adopt ready results at compaction checkpoints while preserving newer messages, tool outputs, and world-state updates. Continue normally if an optional result is rejected. - Run one foreground compaction if an adopted result still exceeds the context limit. - Prevent cancelled turns from launching new background work while preserving already-running work for later reuse. ## Testing Add unit and scenario coverage for stale history and account changes, cancellation, pending-result reuse and fallback, retained tool outputs, compaction event pairing, and token usage accounting. Add a context snapshot covering a collaboration-mode change while compaction runs. GitOrigin-RevId: 04ac028d9fe2e3ea64aacaaeb351c3c776cf3940
This commit is contained in:
@@ -1479,7 +1479,6 @@ impl<'a> ResponsesHttpExecutor<'a> {
|
||||
}
|
||||
|
||||
impl ModelClientSession {
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn turn_state_with_auth_owner(&self) -> (Option<String>, Option<u64>) {
|
||||
let owner = self.client.auth_owner_generation();
|
||||
let turn_state = (self.websocket_session.auth_owner_generation == owner)
|
||||
|
||||
@@ -388,6 +388,14 @@ impl CodexThread {
|
||||
self.session.flush_rollout().await
|
||||
}
|
||||
|
||||
/// Test barrier: waits for background compaction without consuming its published result.
|
||||
/// Call after the request starts, with foreground work gated to prevent new launches.
|
||||
#[doc(hidden)]
|
||||
pub async fn wait_for_background_compaction_for_test(&self) {
|
||||
crate::session::async_compaction::wait_for_background_compaction_for_test(&self.session)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub async fn submit_with_trace(
|
||||
&self,
|
||||
op: Op,
|
||||
|
||||
@@ -117,7 +117,7 @@ pub(crate) async fn run_inline_remote_auto_compact_task(
|
||||
world_state: Arc<WorldState>,
|
||||
reason: CompactionReason,
|
||||
phase: CompactionPhase,
|
||||
candidate: Option<BackgroundCompactionCandidate>,
|
||||
candidate: Option<Box<BackgroundCompactionCandidate>>,
|
||||
fallback: CompactionFallback,
|
||||
) -> CodexResult<bool> {
|
||||
let compaction_metadata = CompactionTurnMetadata::new(
|
||||
@@ -172,7 +172,7 @@ async fn run_remote_compact_task_inner(
|
||||
client_session: Option<&mut ModelClientSession>,
|
||||
world_state: Arc<WorldState>,
|
||||
compaction_metadata: CompactionTurnMetadata,
|
||||
candidate: Option<BackgroundCompactionCandidate>,
|
||||
candidate: Option<Box<BackgroundCompactionCandidate>>,
|
||||
fallback: CompactionFallback,
|
||||
) -> CodexResult<bool> {
|
||||
let turn_context = &step_context.turn;
|
||||
@@ -231,7 +231,7 @@ async fn run_remote_compact_task_inner(
|
||||
compaction_metadata,
|
||||
&mut analytics_details,
|
||||
started_item.as_ref(),
|
||||
candidate,
|
||||
candidate.map(|candidate| *candidate),
|
||||
fallback,
|
||||
))
|
||||
.await;
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
//! Materializes and installs foreground or prepared V2 compaction through one lifecycle.
|
||||
//! A rejected prepared result may fall back once; building and installation share one path.
|
||||
//! Materializes and installs foreground or prepared V2 compaction through a shared path.
|
||||
//! A rejected or still-over-budget prepared result may fall back to one foreground request.
|
||||
//! Prepared prefixes are validated by the history replacement primitive before mutation.
|
||||
|
||||
use super::*;
|
||||
@@ -9,6 +9,10 @@ use crate::context_manager::updates::merge_world_state_updates;
|
||||
use crate::session::CompactedHistoryInstallation;
|
||||
use codex_rollout_trace::CompactionCheckpointTracePayload;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "compact_remote_v2_installation_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn run_compaction(
|
||||
sess: &Arc<Session>,
|
||||
@@ -18,7 +22,7 @@ pub(super) async fn run_compaction(
|
||||
world_state: &WorldState,
|
||||
compaction_metadata: CompactionTurnMetadata,
|
||||
analytics_details: &mut CompactionAnalyticsDetails,
|
||||
started_item: Option<&TurnItem>,
|
||||
mut started_item: Option<&TurnItem>,
|
||||
candidate: Option<BackgroundCompactionCandidate>,
|
||||
fallback: CompactionFallback,
|
||||
) -> CodexResult<bool> {
|
||||
@@ -35,8 +39,20 @@ pub(super) async fn run_compaction(
|
||||
Some(candidate),
|
||||
)
|
||||
.await?;
|
||||
if installed || fallback == CompactionFallback::Skip {
|
||||
return Ok(installed);
|
||||
if installed {
|
||||
let status = crate::session::context_window::context_window_token_status(
|
||||
sess,
|
||||
&replacement_step_context.turn,
|
||||
)
|
||||
.await;
|
||||
if !status.token_limit_reached {
|
||||
return Ok(true);
|
||||
}
|
||||
// The live suffix can leave an adopted prefix over budget. Try the existing
|
||||
// foreground path once, with a new item: adoption already completed its item.
|
||||
started_item = None;
|
||||
} else if fallback == CompactionFallback::Skip {
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
build_and_install_compaction(
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
//! Stale prefixes are skipped or replaced with a fresh request according to caller intent.
|
||||
|
||||
use super::*;
|
||||
use crate::compact_remote_v2::RemoteCompactV2Attempt;
|
||||
use crate::context::world_state::WorldState;
|
||||
use crate::session::CompactionPrefix;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use core_test_support::responses;
|
||||
use pretty_assertions::assert_eq;
|
||||
use test_case::test_case;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum HistoryChange {
|
||||
Replacement,
|
||||
Account,
|
||||
}
|
||||
|
||||
#[test_case(HistoryChange::Replacement, CompactionFallback::Skip; "replacement optional")]
|
||||
#[test_case(HistoryChange::Account, CompactionFallback::Skip; "account optional")]
|
||||
#[test_case(HistoryChange::Replacement, CompactionFallback::Foreground; "replacement required")]
|
||||
#[test_case(HistoryChange::Account, CompactionFallback::Foreground; "account required")]
|
||||
#[tokio::test]
|
||||
async fn prepared_prefix_rejection_respects_caller_intent(
|
||||
change: HistoryChange,
|
||||
fallback: CompactionFallback,
|
||||
) {
|
||||
let server = responses::start_mock_server().await;
|
||||
let mut provider =
|
||||
codex_model_provider_info::built_in_model_providers(/*openai_base_url*/ None)["openai"]
|
||||
.clone();
|
||||
provider.base_url = Some(format!("{}/v1", server.uri()));
|
||||
provider.supports_websockets = false;
|
||||
let server_uri = server.uri();
|
||||
let (sess, turn, _events) =
|
||||
crate::session::tests::make_session_and_context_with_auth_and_config_and_rx(
|
||||
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
|
||||
Vec::new(),
|
||||
move |config| {
|
||||
config.model = Some("gpt-5.2".to_string());
|
||||
config.model_provider = provider;
|
||||
config.chatgpt_base_url = server_uri;
|
||||
let _ = config.features.disable(Feature::TokenBudget);
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let requests = responses::mount_sse_once(
|
||||
&server,
|
||||
responses::sse(vec![
|
||||
serde_json::json!({"type":"response.output_item.done", "item": {
|
||||
"type":"compaction", "encrypted_content":"fresh-summary"
|
||||
}}),
|
||||
serde_json::json!({"type":"response.completed", "response":{"id":"foreground"}}),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
let step = sess
|
||||
.capture_step_context(turn, &CancellationToken::new())
|
||||
.await
|
||||
.expect("capture step");
|
||||
let original: ResponseItem = serde_json::from_value(serde_json::json!({
|
||||
"type": "message", "role": "user",
|
||||
"content": [{"type": "input_text", "text": "original turn"}]
|
||||
}))
|
||||
.expect("user message");
|
||||
sess.replace_history(vec![original.clone()], /*reference_context_item*/ None)
|
||||
.await;
|
||||
let history = sess.clone_history().await;
|
||||
let auth = sess.services.model_client.auth_owner_generation();
|
||||
let mut candidate = BackgroundCompactionCandidate {
|
||||
prefix: CompactionPrefix {
|
||||
generation: history.history_version(),
|
||||
len: history.annotated_items().len(),
|
||||
auth_owner_generation: auth,
|
||||
},
|
||||
model_info: Arc::clone(&step.settings.model_info),
|
||||
cutoff_world_state: Arc::new(WorldState::default()),
|
||||
attempt: RemoteCompactV2Attempt {
|
||||
input_goal_ids: Default::default(),
|
||||
trace_input_history: None,
|
||||
prompt_input: vec![original.clone()],
|
||||
prompt_input_metadata: vec![None],
|
||||
compaction_output: ResponseItem::Compaction {
|
||||
id: None,
|
||||
encrypted_content: "stale-summary".to_string(),
|
||||
internal_chat_message_metadata_passthrough: None,
|
||||
},
|
||||
compaction_response_id: "prepared".to_string(),
|
||||
token_usage: Some(TokenUsage {
|
||||
output_tokens: 777,
|
||||
cached_input_tokens: 9,
|
||||
cache_write_input_tokens: 5,
|
||||
..Default::default()
|
||||
}),
|
||||
owned_client_session: None,
|
||||
},
|
||||
};
|
||||
match change {
|
||||
HistoryChange::Replacement => {
|
||||
sess.replace_history(
|
||||
vec![original.clone(), original.clone()],
|
||||
/*reference_context_item*/ None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
HistoryChange::Account => {
|
||||
candidate.prefix.auth_owner_generation = Some(auth.unwrap_or_default().wrapping_add(1));
|
||||
}
|
||||
}
|
||||
// Rewritten histories still fit the captured length; length alone cannot reject them.
|
||||
let history = sess.clone_history().await;
|
||||
assert!(history.annotated_items().len() >= candidate.prefix.len);
|
||||
if !matches!(change, HistoryChange::Account) {
|
||||
assert_ne!(history.history_version(), candidate.prefix.generation);
|
||||
}
|
||||
let before = sess
|
||||
.clone_history()
|
||||
.await
|
||||
.raw_items()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let mut analytics = CompactionAnalyticsDetails::default();
|
||||
let installed = run_compaction(
|
||||
&sess,
|
||||
&step,
|
||||
&step,
|
||||
/*client_session*/ None,
|
||||
&WorldState::default(),
|
||||
CompactionTurnMetadata::new(
|
||||
CompactionTrigger::Auto,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionImplementation::ResponsesCompactionV2,
|
||||
CompactionPhase::PreTurn,
|
||||
),
|
||||
&mut analytics,
|
||||
/*started_item*/ None,
|
||||
Some(candidate),
|
||||
fallback,
|
||||
)
|
||||
.await
|
||||
.expect("reject stale result without failing compaction");
|
||||
// The foreground response has no usage; stale candidate statistics must not leak.
|
||||
assert_eq!(
|
||||
(
|
||||
analytics.compaction_summary_tokens,
|
||||
analytics.cached_input_tokens,
|
||||
analytics.cache_write_input_tokens
|
||||
),
|
||||
(None, None, None),
|
||||
);
|
||||
let after = sess
|
||||
.clone_history()
|
||||
.await
|
||||
.raw_items()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
match fallback {
|
||||
CompactionFallback::Skip => {
|
||||
assert!(!installed);
|
||||
assert_eq!(analytics.retained_image_count, None);
|
||||
assert_eq!(requests.requests().len(), 0);
|
||||
assert_eq!(after, before);
|
||||
}
|
||||
CompactionFallback::Foreground => {
|
||||
assert!(installed);
|
||||
assert_eq!(requests.requests().len(), 1);
|
||||
let summaries = after
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
ResponseItem::Compaction {
|
||||
encrypted_content, ..
|
||||
} => Some(encrypted_content.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(summaries, ["fresh-summary"]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -31,7 +31,8 @@ struct CompactionAttempt {
|
||||
#[derive(Default)]
|
||||
pub(crate) struct AsyncCompactionState {
|
||||
last_attempt: Option<CompactionAttempt>,
|
||||
result: Option<oneshot::Receiver<BackgroundCompactionCandidate>>,
|
||||
// Keep prepared request state off the turn stack while checkpoints await adoption.
|
||||
result: Option<oneshot::Receiver<Box<BackgroundCompactionCandidate>>>,
|
||||
tasks: TaskTracker,
|
||||
shutdown_cancellation: CancellationToken,
|
||||
}
|
||||
@@ -53,6 +54,7 @@ pub(super) async fn maybe_start(
|
||||
step_context: &Arc<StepContext>,
|
||||
world_state: &Arc<WorldState>,
|
||||
phase: CompactionPhase,
|
||||
turn_cancellation: &CancellationToken,
|
||||
token_status: &ContextWindowTokenStatus,
|
||||
(turn_state, auth_owner_generation): (Option<String>, Option<u64>),
|
||||
) {
|
||||
@@ -87,7 +89,9 @@ pub(super) async fn maybe_start(
|
||||
compaction_hash: model_info.comp_hash.clone(),
|
||||
auth_owner_generation,
|
||||
};
|
||||
if state.shutting_down
|
||||
// Stop blocks new launches; already-running compaction remains session-owned.
|
||||
if turn_cancellation.is_cancelled()
|
||||
|| state.shutting_down
|
||||
|| state.async_compaction.last_attempt.as_ref() == Some(&attempt)
|
||||
|| state.history.annotated_items().is_empty()
|
||||
|| sess.services.model_client.auth_owner_generation() != auth_owner_generation
|
||||
@@ -124,7 +128,7 @@ pub(super) async fn maybe_start(
|
||||
.await
|
||||
{
|
||||
Ok(candidate) => {
|
||||
let _ = send.send(candidate);
|
||||
let _ = send.send(Box::new(candidate));
|
||||
}
|
||||
Err(error) => {
|
||||
debug!(%error, "discarding unsuccessful background compaction");
|
||||
@@ -140,7 +144,7 @@ pub(super) async fn maybe_start(
|
||||
pub(super) async fn take_ready(
|
||||
sess: &Session,
|
||||
turn_context: &TurnContext,
|
||||
) -> Option<BackgroundCompactionCandidate> {
|
||||
) -> Option<Box<BackgroundCompactionCandidate>> {
|
||||
if !enabled(turn_context) {
|
||||
return None;
|
||||
}
|
||||
@@ -153,7 +157,7 @@ pub(super) async fn take_ready(
|
||||
pub(crate) async fn wait_for_pending(
|
||||
sess: &Session,
|
||||
step_context: &StepContext,
|
||||
) -> Option<BackgroundCompactionCandidate> {
|
||||
) -> Option<Box<BackgroundCompactionCandidate>> {
|
||||
if !enabled(&step_context.turn) {
|
||||
return None;
|
||||
}
|
||||
@@ -196,3 +200,15 @@ pub(super) async fn shutdown(sess: &Session) {
|
||||
// Drain completed usage charging before shutdown finishes.
|
||||
tasks.wait().await;
|
||||
}
|
||||
|
||||
/// Test barrier: callers must gate foreground work so it cannot launch another attempt.
|
||||
pub(crate) async fn wait_for_background_compaction_for_test(sess: &Session) {
|
||||
let tasks = sess.state.lock().await.async_compaction.tasks.clone();
|
||||
// Closing only enables waiting; it neither cancels tasks nor prevents future launches.
|
||||
tasks.close();
|
||||
tasks.wait().await;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "async_compaction_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
//! Cancellation prevents new speculation without cancelling session-owned work.
|
||||
|
||||
use super::*;
|
||||
use codex_analytics::CompactionReason;
|
||||
use codex_analytics::CompactionTrigger;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use core_test_support::responses;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test_case::test_case(CompactionTrigger::Auto; "optional ready adoption")]
|
||||
#[test_case::test_case(CompactionTrigger::Manual; "manual ready reuse")]
|
||||
#[tokio::test]
|
||||
async fn cancelled_turn_skips_launch_but_preserves_existing_work(trigger: CompactionTrigger) {
|
||||
let phase = CompactionPhase::MidTurn;
|
||||
let server = responses::start_mock_server().await;
|
||||
let mut provider =
|
||||
codex_model_provider_info::built_in_model_providers(/*openai_base_url*/ None)["openai"]
|
||||
.clone();
|
||||
provider.base_url = Some(format!("{}/v1", server.uri()));
|
||||
provider.supports_websockets = false;
|
||||
let server_uri = server.uri();
|
||||
let (sess, turn, _events) =
|
||||
crate::session::tests::make_session_and_context_with_auth_and_config_and_rx(
|
||||
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
|
||||
Vec::new(),
|
||||
move |config| {
|
||||
config.model = Some("gpt-5.2".to_string());
|
||||
config.model_context_window = Some(100_000);
|
||||
config.model_auto_compact_token_limit = Some(90_000);
|
||||
config.model_provider = provider;
|
||||
config.chatgpt_base_url = server_uri;
|
||||
config.features.enable(Feature::AsyncCompaction).unwrap();
|
||||
config.features.disable(Feature::TokenBudget).unwrap();
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let request = responses::mount_sse_once(
|
||||
&server,
|
||||
responses::sse(vec![
|
||||
serde_json::json!({"type":"response.output_item.done", "item": {
|
||||
"type":"compaction", "encrypted_content":"prepared-summary"
|
||||
}}),
|
||||
responses::ev_completed("prepared"),
|
||||
]),
|
||||
)
|
||||
.await;
|
||||
let turn_cancellation = CancellationToken::new();
|
||||
let step = sess
|
||||
.capture_step_context(Arc::clone(&turn), &turn_cancellation)
|
||||
.await
|
||||
.expect("capture step");
|
||||
let original: ResponseItem = serde_json::from_value(serde_json::json!({
|
||||
"type": "message", "role": "user",
|
||||
"content": [{"type": "input_text", "text": "original turn"}]
|
||||
}))
|
||||
.expect("user message");
|
||||
sess.replace_history(vec![original], /*reference_context_item*/ None)
|
||||
.await;
|
||||
sess.update_token_usage_info(
|
||||
&turn,
|
||||
Some(&TokenUsage {
|
||||
input_tokens: 70_000,
|
||||
total_tokens: 70_000,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.expect("record usage");
|
||||
let status = super::super::context_window::context_window_token_status(&sess, &turn).await;
|
||||
let world_state = Arc::new(WorldState::default());
|
||||
let cancelled = CancellationToken::new();
|
||||
cancelled.cancel();
|
||||
let auth = sess.services.model_client.auth_owner_generation();
|
||||
maybe_start(
|
||||
&sess,
|
||||
&step,
|
||||
&world_state,
|
||||
phase,
|
||||
&cancelled,
|
||||
&status,
|
||||
(Some("routing".to_string()), auth),
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
sess.state
|
||||
.lock()
|
||||
.await
|
||||
.async_compaction
|
||||
.last_attempt
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(request.requests().len(), 0);
|
||||
|
||||
// The same eligible work can launch on a live turn, then outlive its cancellation.
|
||||
maybe_start(
|
||||
&sess,
|
||||
&step,
|
||||
&world_state,
|
||||
phase,
|
||||
&turn_cancellation,
|
||||
&status,
|
||||
(Some("routing".to_string()), auth),
|
||||
)
|
||||
.await;
|
||||
turn_cancellation.cancel();
|
||||
maybe_start(
|
||||
&sess,
|
||||
&step,
|
||||
&world_state,
|
||||
phase,
|
||||
&turn_cancellation,
|
||||
&status,
|
||||
(Some("routing".to_string()), auth),
|
||||
)
|
||||
.await;
|
||||
// Wait for client publication before exercising either ready-result consumer.
|
||||
wait_for_background_compaction_for_test(&sess).await;
|
||||
let next_step = sess
|
||||
.capture_step_context(Arc::clone(&turn), &CancellationToken::new())
|
||||
.await
|
||||
.expect("capture next step");
|
||||
match trigger {
|
||||
CompactionTrigger::Auto => {
|
||||
let candidate = take_ready(&sess, &turn).await.expect("published candidate");
|
||||
let mut client_session = sess.services.model_client.new_session();
|
||||
assert!(
|
||||
crate::compact_remote_v2::run_inline_remote_auto_compact_task(
|
||||
Arc::clone(&sess),
|
||||
Arc::clone(&next_step),
|
||||
next_step,
|
||||
&mut client_session,
|
||||
world_state,
|
||||
CompactionReason::ContextLimit,
|
||||
phase,
|
||||
Some(candidate),
|
||||
crate::compact_remote_v2::CompactionFallback::Skip,
|
||||
)
|
||||
.await
|
||||
.expect("optional adoption")
|
||||
);
|
||||
}
|
||||
CompactionTrigger::Manual => {
|
||||
crate::compact_remote_v2::run_remote_compact_task(
|
||||
Arc::clone(&sess),
|
||||
next_step,
|
||||
world_state,
|
||||
)
|
||||
.await
|
||||
.expect("manual compaction reuses the completed request");
|
||||
}
|
||||
}
|
||||
let history = sess.clone_history().await;
|
||||
let summaries = history
|
||||
.raw_items()
|
||||
.filter_map(|item| match item {
|
||||
ResponseItem::Compaction {
|
||||
encrypted_content, ..
|
||||
} => Some(encrypted_content.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(summaries, ["prepared-summary"]);
|
||||
assert_eq!(request.requests().len(), 1);
|
||||
shutdown(&sess).await;
|
||||
}
|
||||
@@ -238,7 +238,6 @@ use codex_protocol::error::Result as CodexResult;
|
||||
#[cfg(test)]
|
||||
use codex_protocol::exec_output::StreamOutput;
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) mod async_compaction;
|
||||
mod code_mode_warning;
|
||||
mod config_refresh;
|
||||
|
||||
@@ -316,7 +316,7 @@ pub(crate) async fn run_turn(
|
||||
}
|
||||
},
|
||||
);
|
||||
record_context?;
|
||||
let mut world_state = record_context?;
|
||||
|
||||
let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins(
|
||||
&sess,
|
||||
@@ -364,7 +364,8 @@ pub(crate) async fn run_turn(
|
||||
/*candidate*/ None,
|
||||
)
|
||||
.await?;
|
||||
sess.record_context_updates_and_set_reference_context_item(first_step_context.as_ref())
|
||||
world_state = sess
|
||||
.record_context_updates_and_set_reference_context_item(first_step_context.as_ref())
|
||||
.await?;
|
||||
crate::guardian::finalize_guardian_input(
|
||||
&sess,
|
||||
@@ -414,6 +415,23 @@ pub(crate) async fn run_turn(
|
||||
.await;
|
||||
}
|
||||
|
||||
if super::async_compaction::enabled(&turn_context) {
|
||||
let token_status = super::context_window::context_window_token_status(
|
||||
sess.as_ref(),
|
||||
turn_context.as_ref(),
|
||||
)
|
||||
.await;
|
||||
super::async_compaction::maybe_start(
|
||||
&sess,
|
||||
&first_step_context,
|
||||
&world_state,
|
||||
CompactionPhase::PreTurn,
|
||||
&cancellation_token,
|
||||
&token_status,
|
||||
client_session.turn_state_with_auth_owner(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
track_turn_resolved_config_analytics(&sess, &first_step_context, &input).await;
|
||||
|
||||
let mut last_agent_message: Option<String> = None;
|
||||
@@ -514,7 +532,8 @@ pub(crate) async fn run_turn(
|
||||
)
|
||||
.await?;
|
||||
|
||||
sess.record_step_world_state_if_changed(step_context.as_ref())
|
||||
world_state = sess
|
||||
.record_step_world_state_if_changed(step_context.as_ref())
|
||||
.await?;
|
||||
|
||||
// Isolated Guardian reviewers deliberately have no skill catalog.
|
||||
@@ -585,6 +604,21 @@ pub(crate) async fn run_turn(
|
||||
let needs_follow_up = model_needs_follow_up || has_pending_input;
|
||||
let token_limit_reached = token_status.token_limit_reached;
|
||||
|
||||
super::async_compaction::maybe_start(
|
||||
&sess,
|
||||
&step_context,
|
||||
&world_state,
|
||||
if needs_follow_up {
|
||||
CompactionPhase::MidTurn
|
||||
} else {
|
||||
CompactionPhase::PostTurn
|
||||
},
|
||||
&cancellation_token,
|
||||
&token_status,
|
||||
client_session.turn_state_with_auth_owner(),
|
||||
)
|
||||
.await;
|
||||
|
||||
trace!(
|
||||
turn_id = %turn_context.sub_id,
|
||||
total_usage_tokens = token_status.active_context_tokens,
|
||||
@@ -617,9 +651,16 @@ pub(crate) async fn run_turn(
|
||||
);
|
||||
}
|
||||
|
||||
let candidate = if needs_follow_up {
|
||||
super::async_compaction::take_ready(&sess, &turn_context).await
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let should_roll_over = auto_compaction_enabled
|
||||
&& needs_follow_up
|
||||
&& (sess.take_new_context_window_request().await || token_limit_reached);
|
||||
&& (sess.take_new_context_window_request().await
|
||||
|| token_limit_reached
|
||||
|| candidate.is_some());
|
||||
if auto_compaction_enabled {
|
||||
let allow_auto_compact_fallback = !should_roll_over && !token_limit_reached;
|
||||
super::token_budget::maybe_record(
|
||||
@@ -633,40 +674,45 @@ pub(crate) async fn run_turn(
|
||||
|
||||
// as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop.
|
||||
if should_roll_over {
|
||||
if let Err(err) = run_auto_compact(
|
||||
let compacted = run_auto_compact(
|
||||
&sess,
|
||||
Arc::clone(&step_context),
|
||||
Arc::clone(&step_context),
|
||||
&mut client_session,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionPhase::MidTurn,
|
||||
/*candidate*/ None,
|
||||
candidate,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if matches!(err.details(), CodexErrorDetails::TurnAborted) {
|
||||
return Err(err);
|
||||
.await;
|
||||
match compacted {
|
||||
Ok(true) => {
|
||||
if run_pending_session_start_hooks(&sess, &turn_context).await {
|
||||
return Ok(None);
|
||||
}
|
||||
// After a signaled step, drain the steer even if compaction
|
||||
// leaves the estimated context above the token limit.
|
||||
can_drain_pending_input = !model_needs_follow_up
|
||||
|| step_context
|
||||
.preempt
|
||||
.as_ref()
|
||||
.is_some_and(CancellationToken::is_cancelled);
|
||||
continue;
|
||||
}
|
||||
Ok(false) => {}
|
||||
Err(err) => {
|
||||
if matches!(err.details(), CodexErrorDetails::TurnAborted) {
|
||||
return Err(err);
|
||||
}
|
||||
let error = err.to_codex_protocol_error();
|
||||
sess.emit_turn_error_lifecycle(
|
||||
turn_context.as_ref(),
|
||||
error.clone(),
|
||||
err.details(),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
let error = err.to_codex_protocol_error();
|
||||
sess.emit_turn_error_lifecycle(
|
||||
turn_context.as_ref(),
|
||||
error.clone(),
|
||||
err.details(),
|
||||
)
|
||||
.await;
|
||||
return Ok(None);
|
||||
}
|
||||
if run_pending_session_start_hooks(&sess, &turn_context).await {
|
||||
return Ok(None);
|
||||
}
|
||||
// After a signaled step, drain the steer even if compaction
|
||||
// leaves the estimated context above the token limit.
|
||||
can_drain_pending_input = !model_needs_follow_up
|
||||
|| step_context
|
||||
.preempt
|
||||
.as_ref()
|
||||
.is_some_and(CancellationToken::is_cancelled);
|
||||
continue;
|
||||
}
|
||||
|
||||
if !needs_follow_up {
|
||||
@@ -732,45 +778,49 @@ pub(crate) async fn run_turn(
|
||||
// Token-budget resets do not summarize, so preserve their existing rollover
|
||||
// policy. Keep summarizing compaction in this task to serialize history updates.
|
||||
let config = &turn_context.config;
|
||||
let should_compact_after_turn = auto_summarization_enabled
|
||||
if auto_summarization_enabled
|
||||
&& config.model_post_turn_compact_threshold_percent > 0
|
||||
&& super::context_window::context_window_token_status(
|
||||
&& !sess.input_queue.has_pending_input(&sess.active_turn).await
|
||||
&& !cancellation_token.is_cancelled()
|
||||
{
|
||||
let token_status = super::context_window::context_window_token_status(
|
||||
sess.as_ref(),
|
||||
turn_context.as_ref(),
|
||||
)
|
||||
.await
|
||||
.turn_end_compaction_threshold_reached
|
||||
&& !sess.input_queue.has_pending_input(&sess.active_turn).await
|
||||
&& !cancellation_token.is_cancelled();
|
||||
if should_compact_after_turn
|
||||
&& let Err(err) = run_auto_compact(
|
||||
&sess,
|
||||
Arc::clone(&step_context),
|
||||
Arc::clone(&step_context),
|
||||
&mut client_session,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionPhase::PostTurn,
|
||||
/*candidate*/ None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
if matches!(
|
||||
err.details(),
|
||||
CodexErrorDetails::Interrupted | CodexErrorDetails::TurnAborted
|
||||
) {
|
||||
return Err(err);
|
||||
}
|
||||
let error = err.to_codex_protocol_error();
|
||||
if matches!(error, CodexErrorInfo::UsageLimitExceeded) {
|
||||
// Preserve the completed answer while stopping automatic work.
|
||||
sess.emit_turn_error_lifecycle(
|
||||
turn_context.as_ref(),
|
||||
error,
|
||||
err.details(),
|
||||
.await;
|
||||
let foreground_due = token_status.turn_end_compaction_threshold_reached;
|
||||
let candidate =
|
||||
super::async_compaction::take_ready(&sess, &turn_context).await;
|
||||
if (foreground_due || candidate.is_some())
|
||||
&& let Err(err) = run_auto_compact(
|
||||
&sess,
|
||||
Arc::clone(&step_context),
|
||||
Arc::clone(&step_context),
|
||||
&mut client_session,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionPhase::PostTurn,
|
||||
candidate,
|
||||
)
|
||||
.await;
|
||||
.await
|
||||
{
|
||||
if matches!(
|
||||
err.details(),
|
||||
CodexErrorDetails::Interrupted | CodexErrorDetails::TurnAborted
|
||||
) {
|
||||
return Err(err);
|
||||
}
|
||||
let error = err.to_codex_protocol_error();
|
||||
if matches!(error, CodexErrorInfo::UsageLimitExceeded) {
|
||||
// Preserve the completed answer while stopping automatic work.
|
||||
sess.emit_turn_error_lifecycle(
|
||||
turn_context.as_ref(),
|
||||
error,
|
||||
err.details(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
warn!(error = %err, "Post-turn compaction failed; preserving the completed turn");
|
||||
}
|
||||
warn!(error = %err, "Post-turn compaction failed; preserving the completed turn");
|
||||
}
|
||||
break;
|
||||
}
|
||||
@@ -1330,8 +1380,9 @@ async fn run_pre_sampling_compact(
|
||||
let token_status =
|
||||
super::context_window::context_window_token_status(sess.as_ref(), turn_context.as_ref())
|
||||
.await;
|
||||
let candidate = super::async_compaction::take_ready(sess, turn_context).await;
|
||||
// Compact if the configured auto-compaction budget or usable context window is exhausted.
|
||||
if token_status.token_limit_reached {
|
||||
if token_status.token_limit_reached || candidate.is_some() {
|
||||
// Pre-turn compaction runs before run_turn creates the normal sampling step.
|
||||
let step_context = sess
|
||||
.capture_step_context(Arc::clone(turn_context), cancellation_token)
|
||||
@@ -1343,7 +1394,7 @@ async fn run_pre_sampling_compact(
|
||||
client_session,
|
||||
CompactionReason::ContextLimit,
|
||||
CompactionPhase::PreTurn,
|
||||
/*candidate*/ None,
|
||||
candidate,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
@@ -1460,7 +1511,7 @@ async fn run_auto_compact(
|
||||
client_session: &mut ModelClientSession,
|
||||
reason: CompactionReason,
|
||||
phase: CompactionPhase,
|
||||
candidate: Option<BackgroundCompactionCandidate>,
|
||||
candidate: Option<Box<BackgroundCompactionCandidate>>,
|
||||
) -> CodexResult<bool> {
|
||||
let turn_context = &step_context.turn;
|
||||
if matches!(phase, CompactionPhase::PreTurn)
|
||||
|
||||
@@ -11,6 +11,7 @@ use std::time::Duration;
|
||||
use anyhow::Result;
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
use codex_analytics::CompactionTrigger;
|
||||
use codex_config::ConfigLayerSource;
|
||||
use codex_config::ConfigLayerStack;
|
||||
use codex_config::types::McpServerConfig;
|
||||
@@ -18,6 +19,7 @@ use codex_context_fragments::AnsweredQuestion;
|
||||
use codex_context_fragments::ContextualUserFragment;
|
||||
use codex_core::ForkSnapshot;
|
||||
use codex_core::StartThreadOptions;
|
||||
use codex_core::TurnInput;
|
||||
use codex_core::TurnInputRequest;
|
||||
use codex_core::compact::AutoCompactionDisabled;
|
||||
use codex_core::config::Config;
|
||||
@@ -29,6 +31,7 @@ use codex_extension_api::ExtensionDataInit;
|
||||
use codex_extension_api::ExtensionRegistry;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_features::Feature;
|
||||
use codex_history::RolloutItem;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_models_manager::bundled_models_response;
|
||||
use codex_protocol::capabilities::CapabilityRootLocation;
|
||||
@@ -80,6 +83,7 @@ use core_test_support::skip_if_wine_exec;
|
||||
use core_test_support::stdio_server_bin;
|
||||
use core_test_support::streaming_sse::StreamingSseChunk;
|
||||
use core_test_support::streaming_sse::start_streaming_sse_server;
|
||||
use core_test_support::test_codex::TestCodex;
|
||||
use core_test_support::test_codex::environment_config_for_selection;
|
||||
use core_test_support::test_codex::executor_path_uri;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
@@ -2173,3 +2177,601 @@ async fn model_catalog_refresh_preserves_tools_and_history() -> Result<()> {
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn compaction_chunk(
|
||||
gate: Option<oneshot::Receiver<()>>,
|
||||
events: Vec<serde_json::Value>,
|
||||
) -> StreamingSseChunk {
|
||||
StreamingSseChunk {
|
||||
gate,
|
||||
body: sse(events),
|
||||
}
|
||||
}
|
||||
|
||||
fn compaction_output(value: &str) -> serde_json::Value {
|
||||
json!({"type": "response.output_item.done", "item": {
|
||||
"type": "compaction", "encrypted_content": value
|
||||
}})
|
||||
}
|
||||
|
||||
async fn async_compaction_test(
|
||||
responses_uri: &str,
|
||||
features: Vec<Feature>,
|
||||
configure: impl Fn(&mut Config) + Send + Sync + 'static,
|
||||
) -> Result<(wiremock::MockServer, TestCodex)> {
|
||||
let config_server = start_mock_server().await;
|
||||
let base_url = format!("{responses_uri}/v1");
|
||||
let test = test_codex()
|
||||
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
|
||||
.with_config(move |config| {
|
||||
configure_scenario_catalog(config);
|
||||
config.model_provider.base_url = Some(base_url);
|
||||
config.model_context_window = Some(100_000);
|
||||
for feature in features {
|
||||
config
|
||||
.features
|
||||
.enable(feature)
|
||||
.expect("enable scenario feature");
|
||||
}
|
||||
config
|
||||
.features
|
||||
.enable(Feature::AsyncCompaction)
|
||||
.expect("enable async compaction");
|
||||
config
|
||||
.features
|
||||
.disable(Feature::EnableRequestCompression)
|
||||
.expect("capture raw requests");
|
||||
configure(config);
|
||||
})
|
||||
.with_model("gpt-6-astra")
|
||||
.build_with_auto_env(&config_server)
|
||||
.await?;
|
||||
Ok((config_server, test))
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum LiveSuffix {
|
||||
User,
|
||||
Tool,
|
||||
RetainedTool,
|
||||
OverBudgetRetainedTool,
|
||||
}
|
||||
|
||||
#[test_case::test_case(LiveSuffix::User; "newer user")]
|
||||
#[test_case::test_case(LiveSuffix::Tool; "only assistant and tools")]
|
||||
#[test_case::test_case(LiveSuffix::RetainedTool; "retained standalone output")]
|
||||
#[test_case::test_case(LiveSuffix::OverBudgetRetainedTool; "ready adoption still over budget")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn astra_async_compaction_keeps_work_added_after_its_snapshot(
|
||||
suffix: LiveSuffix,
|
||||
) -> Result<()> {
|
||||
use codex_protocol::config_types::CollaborationMode;
|
||||
use codex_protocol::config_types::ModeKind;
|
||||
use codex_protocol::config_types::Settings;
|
||||
use codex_protocol::protocol::ThreadSettingsOverrides;
|
||||
use core_test_support::responses::ev_completed_with_tokens;
|
||||
|
||||
skip_if_no_network!(Ok(()));
|
||||
let over_budget = matches!(suffix, LiveSuffix::OverBudgetRetainedTool);
|
||||
let retained_output = if over_budget {
|
||||
format!("delegated work {}", "x ".repeat(12_000))
|
||||
} else {
|
||||
"delegated work".to_string()
|
||||
};
|
||||
let (release_background, background_gate) = oneshot::channel();
|
||||
let (release_foreground, foreground_gate) = oneshot::channel();
|
||||
let mut responses = vec![
|
||||
vec![compaction_chunk(
|
||||
/*gate*/ None,
|
||||
vec![
|
||||
ev_assistant_message("draft", "I drafted the launch notes."),
|
||||
ev_completed_with_tokens(
|
||||
"draft-response",
|
||||
/*total_tokens*/ if over_budget { 2_000 } else { 70_000 },
|
||||
),
|
||||
],
|
||||
)],
|
||||
vec![compaction_chunk(
|
||||
Some(background_gate),
|
||||
vec![
|
||||
compaction_output("ASYNC_LAUNCH_CHECKPOINT"),
|
||||
ev_completed_with_tokens("compact-response", /*total_tokens*/ 3_000),
|
||||
],
|
||||
)],
|
||||
vec![
|
||||
compaction_chunk(
|
||||
/*gate*/ None,
|
||||
vec![
|
||||
ev_response_created("team-response"),
|
||||
ev_assistant_message("team-check", "I will check who is available."),
|
||||
],
|
||||
),
|
||||
compaction_chunk(
|
||||
Some(foreground_gate),
|
||||
vec![
|
||||
ev_function_call_with_namespace("team", "collaboration", "list_agents", "{}"),
|
||||
ev_completed_with_tokens(
|
||||
"team-response",
|
||||
/*total_tokens*/ if over_budget { 9_000 } else { 75_000 },
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
vec![compaction_chunk(
|
||||
/*gate*/ None,
|
||||
vec![
|
||||
ev_assistant_message("answer", "I can finish the notes myself."),
|
||||
ev_completed_with_tokens(
|
||||
"answer-response",
|
||||
/*total_tokens*/ if over_budget { 10 } else { 3_000 },
|
||||
),
|
||||
],
|
||||
)],
|
||||
];
|
||||
if over_budget {
|
||||
responses.insert(
|
||||
3,
|
||||
vec![compaction_chunk(
|
||||
/*gate*/ None,
|
||||
vec![
|
||||
compaction_output("SYNC"),
|
||||
ev_completed_with_tokens("sync", /*total_tokens*/ 3_000),
|
||||
],
|
||||
)],
|
||||
);
|
||||
}
|
||||
let (streaming, _completions) = start_streaming_sse_server(responses).await;
|
||||
let features = match suffix {
|
||||
LiveSuffix::User => vec![Feature::IncrementalTools],
|
||||
LiveSuffix::Tool => vec![],
|
||||
LiveSuffix::RetainedTool | LiveSuffix::OverBudgetRetainedTool => {
|
||||
vec![Feature::RetainClientToolOutputs]
|
||||
}
|
||||
};
|
||||
let (_config_server, test) = async_compaction_test(streaming.uri(), features, move |config| {
|
||||
if over_budget {
|
||||
config.async_compaction_threshold_percent = 1;
|
||||
config.model_auto_compact_token_limit = Some(5_000);
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
test.codex.ensure_rollout_materialized().await;
|
||||
let compaction_event = |event: &EventMsg| match event {
|
||||
EventMsg::ItemStarted(event) if matches!(event.item, TurnItem::ContextCompaction(_)) => {
|
||||
Some(("started", event.item.id()))
|
||||
}
|
||||
EventMsg::ItemCompleted(event) if matches!(event.item, TurnItem::ContextCompaction(_)) => {
|
||||
Some(("completed", event.item.id()))
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
test.submit_text_turn("Draft the launch notes.").await?;
|
||||
let timeout = Duration::from_secs(/*secs*/ 20);
|
||||
tokio::time::timeout(timeout, streaming.wait_for_request_count(/*count*/ 2)).await?;
|
||||
let next_input = if matches!(suffix, LiveSuffix::User) {
|
||||
TurnInputRequest::user_input(vec![text("Check who can finish the notes.")])
|
||||
} else {
|
||||
let item = serde_json::from_value(json!({
|
||||
"type": "function_call_output",
|
||||
"name": "send_message_to_thread",
|
||||
"namespace": "codex_app",
|
||||
"output": retained_output,
|
||||
}))?;
|
||||
TurnInputRequest::new(
|
||||
if matches!(
|
||||
suffix,
|
||||
LiveSuffix::RetainedTool | LiveSuffix::OverBudgetRetainedTool
|
||||
) {
|
||||
TurnInput::AnnotatedResponseItem(
|
||||
codex_protocol::turn_input::AnnotatedResponseItem {
|
||||
item,
|
||||
annotations: codex_protocol::turn_input::ResponseItemAnnotations {
|
||||
retain: true,
|
||||
},
|
||||
},
|
||||
)
|
||||
} else {
|
||||
TurnInput::ResponseItem(item)
|
||||
},
|
||||
)
|
||||
};
|
||||
test.codex
|
||||
.start_or_steer_turn(next_input.with_thread_settings(ThreadSettingsOverrides {
|
||||
collaboration_mode: Some(CollaborationMode {
|
||||
mode: ModeKind::Plan,
|
||||
settings: Settings {
|
||||
model: "gpt-6-astra".to_string(),
|
||||
reasoning_effort: None,
|
||||
developer_instructions: Some(
|
||||
"Plan Mode: prepare a plan before making changes.".to_string(),
|
||||
),
|
||||
},
|
||||
}),
|
||||
..Default::default()
|
||||
}))
|
||||
.await?;
|
||||
tokio::time::timeout(timeout, streaming.wait_for_request_count(/*count*/ 3)).await?;
|
||||
wait_for_event(&test.codex, |event| {
|
||||
assert_eq!(compaction_event(event), None);
|
||||
matches!(event, EventMsg::ItemCompleted(event)
|
||||
if matches!(&event.item, TurnItem::AgentMessage(message) if message.id == "team-check"))
|
||||
})
|
||||
.await;
|
||||
release_background
|
||||
.send(())
|
||||
.expect("release background compaction");
|
||||
// Wait for client publication while foreground work is gated, keeping adoption optional.
|
||||
tokio::time::timeout(
|
||||
timeout,
|
||||
test.codex.wait_for_background_compaction_for_test(),
|
||||
)
|
||||
.await?;
|
||||
release_foreground
|
||||
.send(())
|
||||
.expect("release foreground work");
|
||||
let mut compaction_events = Vec::new();
|
||||
wait_for_event(&test.codex, |event| {
|
||||
compaction_events.extend(compaction_event(event));
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
tokio::time::timeout(
|
||||
timeout,
|
||||
test.codex.wait_for_background_compaction_for_test(),
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(compaction_events.len(), if over_budget { 4 } else { 2 });
|
||||
for pair in compaction_events.chunks_exact(2) {
|
||||
assert_eq!((pair[0].0, pair[1].0), ("started", "completed"));
|
||||
assert_eq!(pair[0].1, pair[1].1);
|
||||
}
|
||||
if over_budget {
|
||||
assert_ne!(compaction_events[0].1, compaction_events[2].1);
|
||||
}
|
||||
|
||||
let requests = streaming
|
||||
.requests()
|
||||
.await
|
||||
.iter()
|
||||
.map(|body| serde_json::from_slice::<serde_json::Value>(body))
|
||||
.collect::<serde_json::Result<Vec<_>>>()?;
|
||||
assert_eq!(requests.len(), if over_budget { 5 } else { 4 });
|
||||
if over_budget {
|
||||
let continuation = requests[4]["input"].as_array().expect("continuation input");
|
||||
assert!(
|
||||
continuation
|
||||
.iter()
|
||||
.any(|item| item["encrypted_content"] == "SYNC")
|
||||
);
|
||||
assert!(
|
||||
continuation
|
||||
.iter()
|
||||
.any(|item| item["output"] == retained_output)
|
||||
);
|
||||
}
|
||||
let history = test
|
||||
.codex
|
||||
.load_history(/*include_archived*/ false)
|
||||
.await?
|
||||
.items;
|
||||
assert!(history.windows(2).any(|items| matches!(items,
|
||||
[RolloutItem::Compacted(_), RolloutItem::WorldState(baseline)]
|
||||
if baseline.full && baseline.state["collaboration_mode"]["mode"] == "plan"
|
||||
)));
|
||||
let input = requests[3]["input"].as_array().expect("request input");
|
||||
// The full snapshot normalizes mode text; check the cutoff and later update explicitly.
|
||||
let position = |text: &str| {
|
||||
let positions = input
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(i, item)| item.to_string().contains(text).then_some(i))
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(positions.len(), 1, "{text}");
|
||||
positions[0]
|
||||
};
|
||||
let default = position("# Collaboration Mode: Default");
|
||||
let draft = position("Draft the launch notes.");
|
||||
let compaction = position("ASYNC_LAUNCH_CHECKPOINT");
|
||||
let plan = position("Plan Mode");
|
||||
if matches!(suffix, LiveSuffix::User) {
|
||||
let user = position("Check who can finish the notes.");
|
||||
assert!(draft < compaction && compaction < default && default < plan && plan < user);
|
||||
let entries = requests.iter().map(SnapshotEntry::body).collect::<Vec<_>>();
|
||||
insta::assert_snapshot!(
|
||||
"astra_async_compaction_with_live_suffix",
|
||||
context_snapshot::format_context_snapshot(
|
||||
"Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode while it runs. Adoption restores Default mode at the cutoff, then preserves the Plan update, newer user message, and assistant/tool exchange in their original order.",
|
||||
&entries,
|
||||
&ContextSnapshotOptions::default().rewrite_known_segments(),
|
||||
)
|
||||
);
|
||||
} else {
|
||||
let tool = position("delegated work");
|
||||
if matches!(
|
||||
suffix,
|
||||
LiveSuffix::RetainedTool | LiveSuffix::OverBudgetRetainedTool
|
||||
) {
|
||||
assert!(draft < compaction && compaction < default && default < plan && plan < tool);
|
||||
} else {
|
||||
assert!(default < draft && draft < compaction && compaction < plan && plan < tool);
|
||||
}
|
||||
let assistant = position("I will check who is available.");
|
||||
let call = input
|
||||
.iter()
|
||||
.position(|item| item["type"] == "function_call" && item["call_id"] == "team")
|
||||
.expect("retained tool call");
|
||||
let output = input
|
||||
.iter()
|
||||
.position(|item| item["type"] == "function_call_output" && item["call_id"] == "team")
|
||||
.expect("retained tool output");
|
||||
assert!(tool < assistant && assistant < call && call < output);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(PartialEq)]
|
||||
enum PendingCompactionOutcome {
|
||||
Absent,
|
||||
Success,
|
||||
SuccessOverBudget,
|
||||
StaleSuccess,
|
||||
Failure,
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
#[test_case::test_case(PendingCompactionOutcome::Success, CompactionTrigger::Auto; "auto uses pending result")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::Failure, CompactionTrigger::Auto; "auto falls back after failure")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::StaleSuccess, CompactionTrigger::Auto; "auto falls back after account changes")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::Interrupted, CompactionTrigger::Auto; "auto can be interrupted")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::Absent, CompactionTrigger::Manual; "manual without background")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::Success, CompactionTrigger::Manual; "manual uses pending result")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::Failure, CompactionTrigger::Manual; "manual falls back after failure")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::StaleSuccess, CompactionTrigger::Manual; "manual falls back after account changes")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::Interrupted, CompactionTrigger::Manual; "manual can be interrupted")]
|
||||
#[test_case::test_case(PendingCompactionOutcome::SuccessOverBudget, CompactionTrigger::Auto; "pending adoption still over budget")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn async_compaction_reuses_pending_work(
|
||||
outcome: PendingCompactionOutcome,
|
||||
trigger: CompactionTrigger,
|
||||
) -> Result<()> {
|
||||
use core_test_support::responses::ev_completed_with_tokens;
|
||||
|
||||
skip_if_no_network!(Ok(()));
|
||||
let over_budget = outcome == PendingCompactionOutcome::SuccessOverBudget;
|
||||
let retained_output = format!("LATE_OUTPUT {}", "x ".repeat(12_000));
|
||||
let manual = matches!(trigger, CompactionTrigger::Manual);
|
||||
let has_background = outcome != PendingCompactionOutcome::Absent;
|
||||
let seed_tokens = if over_budget {
|
||||
2_000
|
||||
} else if has_background {
|
||||
70_000
|
||||
} else {
|
||||
10_000
|
||||
};
|
||||
let valid_output = outcome != PendingCompactionOutcome::Failure;
|
||||
let succeeds = !matches!(
|
||||
outcome,
|
||||
PendingCompactionOutcome::Absent
|
||||
| PendingCompactionOutcome::SuccessOverBudget
|
||||
| PendingCompactionOutcome::Failure
|
||||
| PendingCompactionOutcome::StaleSuccess
|
||||
);
|
||||
let response = |id, event, tokens| vec![event, ev_completed_with_tokens(id, tokens)];
|
||||
let immediate = |events| vec![compaction_chunk(/*gate*/ None, events)];
|
||||
let (release_background, background_gate) = oneshot::channel();
|
||||
let mut responses = vec![immediate(response(
|
||||
"seed",
|
||||
ev_assistant_message("seed", "seed"),
|
||||
seed_tokens,
|
||||
))];
|
||||
if has_background {
|
||||
responses.push(vec![compaction_chunk(
|
||||
Some(background_gate),
|
||||
response(
|
||||
"background",
|
||||
if valid_output {
|
||||
compaction_output("ASYNC")
|
||||
} else {
|
||||
ev_assistant_message("invalid", "missing compaction output")
|
||||
},
|
||||
/*tokens*/ 3_000,
|
||||
),
|
||||
)]);
|
||||
}
|
||||
if !manual {
|
||||
responses.push(immediate(response(
|
||||
"second",
|
||||
ev_function_call_with_namespace("team", "collaboration", "list_agents", "{}"),
|
||||
/*tokens*/ if over_budget { 9_000 } else { 91_000 },
|
||||
)));
|
||||
}
|
||||
if !succeeds {
|
||||
responses.push(immediate(response(
|
||||
"sync",
|
||||
compaction_output("SYNC"),
|
||||
/*tokens*/ 3_000,
|
||||
)));
|
||||
}
|
||||
responses.push(immediate(response(
|
||||
"continuation",
|
||||
ev_assistant_message("continuation", "continuation"),
|
||||
/*tokens*/ if over_budget { 10 } else { 3_000 },
|
||||
)));
|
||||
let (streaming, mut completions) = start_streaming_sse_server(responses).await;
|
||||
let (_config_server, test) = async_compaction_test(
|
||||
streaming.uri(),
|
||||
if over_budget {
|
||||
vec![Feature::RetainClientToolOutputs]
|
||||
} else {
|
||||
vec![]
|
||||
},
|
||||
move |config| {
|
||||
if over_budget {
|
||||
config.async_compaction_threshold_percent = 1;
|
||||
config.model_auto_compact_token_limit = Some(5_000);
|
||||
}
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
test.codex.ensure_rollout_materialized().await;
|
||||
let timeout = Duration::from_secs(/*secs*/ 20);
|
||||
test.submit_text_turn("first").await?;
|
||||
if has_background {
|
||||
tokio::time::timeout(timeout, streaming.wait_for_request_count(/*count*/ 2)).await?;
|
||||
}
|
||||
if manual {
|
||||
test.codex.submit(Op::Compact).await?;
|
||||
} else if over_budget {
|
||||
test.codex.start_or_steer_turn(TurnInputRequest::new(TurnInput::AnnotatedResponseItem(
|
||||
codex_protocol::turn_input::AnnotatedResponseItem {
|
||||
item: serde_json::from_value(json!({"type": "function_call_output", "name": "send_message_to_thread", "namespace": "codex_app", "output": retained_output}))?,
|
||||
annotations: codex_protocol::turn_input::ResponseItemAnnotations { retain: true },
|
||||
},
|
||||
))).await?;
|
||||
} else {
|
||||
test.codex
|
||||
.start_or_steer_turn(TurnInputRequest::user_input(vec![text("newer")]))
|
||||
.await?;
|
||||
}
|
||||
let started = wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::ItemStarted(event)
|
||||
if matches!(event.item, TurnItem::ContextCompaction(_)))
|
||||
})
|
||||
.await;
|
||||
if has_background {
|
||||
assert_eq!(streaming.requests().await.len(), if manual { 2 } else { 3 });
|
||||
}
|
||||
if outcome == PendingCompactionOutcome::Interrupted {
|
||||
test.codex.submit(Op::Interrupt).await?;
|
||||
wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnAborted(_))
|
||||
})
|
||||
.await;
|
||||
release_background.send(()).expect("release unused result");
|
||||
completions.remove(1).await?;
|
||||
test.codex.shutdown_and_wait().await?;
|
||||
assert_eq!(streaming.requests().await.len(), if manual { 2 } else { 3 });
|
||||
return Ok(());
|
||||
}
|
||||
if outcome == PendingCompactionOutcome::StaleSuccess {
|
||||
let auth = test.thread_manager.auth_manager();
|
||||
let mut tokens = auth.auth_cached().expect("test auth").get_token_data()?;
|
||||
tokens.id_token.raw_jwt = "e30.e30.signature".to_string();
|
||||
tokens.account_id = Some("second-account".to_string());
|
||||
tokens.access_token = "second-account-token".to_string();
|
||||
std::fs::write(
|
||||
test.home.path().join("auth.json"),
|
||||
serde_json::to_vec(
|
||||
&json!({"auth_mode":"chatgpt", "tokens":tokens, "last_refresh":chrono::Utc::now()}),
|
||||
)?,
|
||||
)?;
|
||||
auth.reload().await;
|
||||
}
|
||||
if has_background {
|
||||
release_background.send(()).expect("release pending result");
|
||||
}
|
||||
let EventMsg::ItemStarted(started) = started else {
|
||||
unreachable!();
|
||||
};
|
||||
let mut starts = vec![started.item.id()];
|
||||
let mut completed = Vec::new();
|
||||
wait_for_event(&test.codex, |event| {
|
||||
if let EventMsg::ItemStarted(event) = event
|
||||
&& matches!(event.item, TurnItem::ContextCompaction(_))
|
||||
{
|
||||
starts.push(event.item.id());
|
||||
}
|
||||
if let EventMsg::ItemCompleted(event) = event
|
||||
&& matches!(event.item, TurnItem::ContextCompaction(_))
|
||||
{
|
||||
completed.push(event.item.id());
|
||||
}
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
assert_eq!(completed, starts);
|
||||
assert_eq!(starts.len(), if over_budget { 2 } else { 1 });
|
||||
assert_eq!(starts.iter().collect::<HashSet<_>>().len(), starts.len());
|
||||
tokio::time::timeout(
|
||||
timeout,
|
||||
test.codex.wait_for_background_compaction_for_test(),
|
||||
)
|
||||
.await?;
|
||||
if manual {
|
||||
test.submit_text_turn("after manual compaction").await?;
|
||||
}
|
||||
|
||||
let history = test.codex.load_history(/*include_archived*/ false).await?;
|
||||
let usage = |id| {
|
||||
history
|
||||
.items
|
||||
.iter()
|
||||
.find_map(|item| match item {
|
||||
RolloutItem::TokenUsageRecord(record) if record.response_id == id => Some(record),
|
||||
_ => None,
|
||||
})
|
||||
.expect("persisted usage")
|
||||
};
|
||||
assert!(history.items.iter().all(|item| !matches!(item,
|
||||
RolloutItem::TokenUsageRecord(record) if record.response_id == "background"
|
||||
)));
|
||||
let continuation = usage("continuation");
|
||||
assert_eq!(
|
||||
(
|
||||
continuation.turn_token_usage.total_tokens,
|
||||
continuation.thread_token_usage.total_tokens,
|
||||
),
|
||||
if over_budget {
|
||||
(12_010, 14_010)
|
||||
} else if manual {
|
||||
(3_000, seed_tokens + if succeeds { 3_000 } else { 6_000 })
|
||||
} else if succeeds {
|
||||
(94_000, 164_000)
|
||||
} else {
|
||||
(97_000, 167_000)
|
||||
},
|
||||
);
|
||||
|
||||
let requests = streaming
|
||||
.requests()
|
||||
.await
|
||||
.iter()
|
||||
.map(|body| serde_json::from_slice::<serde_json::Value>(body))
|
||||
.collect::<serde_json::Result<Vec<_>>>()?;
|
||||
assert_eq!(
|
||||
requests.len(),
|
||||
(if manual { 2 } else { 3 }) + usize::from(has_background) + usize::from(!succeeds)
|
||||
);
|
||||
if over_budget {
|
||||
for request in &requests[3..] {
|
||||
assert!(
|
||||
request["input"]
|
||||
.as_array()
|
||||
.expect("compaction or continuation input")
|
||||
.iter()
|
||||
.any(|item| item["output"] == retained_output)
|
||||
);
|
||||
}
|
||||
}
|
||||
if manual && !succeeds {
|
||||
let foreground = &requests[requests.len() - 2];
|
||||
let metadata: serde_json::Value = serde_json::from_str(
|
||||
foreground["client_metadata"]["x-codex-turn-metadata"]
|
||||
.as_str()
|
||||
.expect("manual foreground metadata"),
|
||||
)?;
|
||||
assert_eq!(metadata["compaction"]["trigger"], "manual");
|
||||
assert_eq!(metadata["compaction"]["reason"], "user_requested");
|
||||
assert_eq!(metadata["compaction"]["phase"], "standalone_turn");
|
||||
}
|
||||
let summaries: Vec<_> = requests.last().expect("continuation")["input"]
|
||||
.as_array()
|
||||
.expect("foreground input")
|
||||
.iter()
|
||||
.filter(|item| item["type"] == "compaction")
|
||||
.map(|item| &item["encrypted_content"])
|
||||
.collect();
|
||||
assert_eq!(summaries, [&json!(if succeeds { "ASYNC" } else { "SYNC" })]);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
+117
@@ -0,0 +1,117 @@
|
||||
---
|
||||
source: core/tests/suite/scenarios.rs
|
||||
expression: "context_snapshot::format_context_snapshot(\"Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode while it runs. Adoption restores Default mode at the cutoff, then preserves the Plan update, newer user message, and assistant/tool exchange in their original order.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)"
|
||||
---
|
||||
Scenario: Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode while it runs. Adoption restores Default mode at the cutoff, then preserves the Plan update, newer user message, and assistant/tool exchange in their original order.
|
||||
|
||||
## Window 1
|
||||
-- request 1 (request) --
|
||||
00:additional_tools/developer (3; hash=bf2e30605ff5dc17):
|
||||
- namespace/functions; hash=fada4a476a3754dd
|
||||
- custom/exec
|
||||
- function/wait
|
||||
- function/request_user_input
|
||||
- function/request_user_input_async
|
||||
- namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317
|
||||
- function/sleep
|
||||
- namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63
|
||||
- function/followup_task
|
||||
- function/interrupt_agent
|
||||
- function/list_agents
|
||||
- function/send_message
|
||||
- function/spawn_agent
|
||||
- function/wait_agent
|
||||
01:message/developer:
|
||||
<MODEL_INSTRUCTIONS>
|
||||
02:message/developer[2]:
|
||||
[01] <PERMISSIONS_INSTRUCTIONS>
|
||||
[02] <COLLABORATION_MODE>
|
||||
03:message/developer:
|
||||
<MULTI_AGENT_ROLE>
|
||||
04:message/developer:
|
||||
<MULTI_AGENT_MODE>
|
||||
05:message/user:
|
||||
<ENVIRONMENT_CONTEXT>
|
||||
06:message/user:
|
||||
Draft the launch notes.
|
||||
-- request 2 (request) --
|
||||
07:message/assistant:
|
||||
I drafted the launch notes.
|
||||
08:compaction_trigger
|
||||
|
||||
## Window 2 (after request 2: input diverged at item 08)
|
||||
-- request 3 (request) --
|
||||
00:additional_tools/developer (3; hash=bf2e30605ff5dc17):
|
||||
- namespace/functions; hash=fada4a476a3754dd
|
||||
- custom/exec
|
||||
- function/wait
|
||||
- function/request_user_input
|
||||
- function/request_user_input_async
|
||||
- namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317
|
||||
- function/sleep
|
||||
- namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63
|
||||
- function/followup_task
|
||||
- function/interrupt_agent
|
||||
- function/list_agents
|
||||
- function/send_message
|
||||
- function/spawn_agent
|
||||
- function/wait_agent
|
||||
01:message/developer:
|
||||
<MODEL_INSTRUCTIONS>
|
||||
02:message/developer[2]:
|
||||
[01] <PERMISSIONS_INSTRUCTIONS>
|
||||
[02] <COLLABORATION_MODE>
|
||||
03:message/developer:
|
||||
<MULTI_AGENT_ROLE>
|
||||
04:message/developer:
|
||||
<MULTI_AGENT_MODE>
|
||||
05:message/user:
|
||||
<ENVIRONMENT_CONTEXT>
|
||||
06:message/user:
|
||||
Draft the launch notes.
|
||||
07:message/assistant:
|
||||
I drafted the launch notes.
|
||||
08:message/developer:
|
||||
<COLLABORATION_MODE>
|
||||
09:message/user:
|
||||
Check who can finish the notes.
|
||||
|
||||
## Window 3 (after request 3: input diverged at item 02)
|
||||
-- request 4 (request) --
|
||||
00:additional_tools/developer (3; hash=bf2e30605ff5dc17):
|
||||
- namespace/functions; hash=fada4a476a3754dd
|
||||
- custom/exec
|
||||
- function/wait
|
||||
- function/request_user_input
|
||||
- function/request_user_input_async
|
||||
- namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317
|
||||
- function/sleep
|
||||
- namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63
|
||||
- function/followup_task
|
||||
- function/interrupt_agent
|
||||
- function/list_agents
|
||||
- function/send_message
|
||||
- function/spawn_agent
|
||||
- function/wait_agent
|
||||
01:message/developer:
|
||||
<MODEL_INSTRUCTIONS>
|
||||
02:message/user:
|
||||
Draft the launch notes.
|
||||
03:compaction:encrypted=true; chars=23; hash=159cfcb231cc70f0
|
||||
04:message/developer[2]:
|
||||
[01] <PERMISSIONS_INSTRUCTIONS>
|
||||
[02] <COLLABORATION_MODE>
|
||||
05:message/developer:
|
||||
<MULTI_AGENT_ROLE>
|
||||
06:message/developer:
|
||||
<MULTI_AGENT_MODE>
|
||||
07:message/user:
|
||||
<ENVIRONMENT_CONTEXT>
|
||||
08:message/developer:
|
||||
<COLLABORATION_MODE>
|
||||
09:message/user:
|
||||
Check who can finish the notes.
|
||||
10:message/assistant:
|
||||
I will check who is available.
|
||||
11:function_call/collaboration.list_agents:{}
|
||||
12:function_call_output:{"agents":[{"agent_name":"/root","agent_status":"running"}]}
|
||||
Reference in New Issue
Block a user