openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

core: reset context for token budget compaction

Browse Source
This commit is contained in:
Michael Bolin
2026-06-23 14:55:02 -07:00
parent 2e69966cd8
commit 1d47a07447
9 changed files with 830 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/compact_token_budget.rs
+87
View File
@@ -0,0 +1,87 @@
use std::sync::Arc;
use crate::context::world_state::WorldState;
use crate::hook_runtime::PostCompactHookOutcome;
use crate::hook_runtime::PreCompactHookOutcome;
use crate::hook_runtime::run_post_compact_hooks;
use crate::hook_runtime::run_pre_compact_hooks;
use crate::session::NewContextWindowMode;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use codex_analytics::CompactionTrigger;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::items::ContextCompactionItem;
use codex_protocol::items::TurnItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TurnStartedEvent;
/// Runs token-budget manual compaction as a normal compaction lifecycle.
///
/// Token-budget compaction skips model/server summarization and installs a fresh context window
/// instead. It is still modeled as compaction so compact hooks and `ContextCompaction` turn items
/// observe the same lifecycle as local or remote compaction.
pub(crate) async fn run_manual_compact_task(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
) -> CodexResult<()> {
let start_event = EventMsg::TurnStarted(TurnStartedEvent {
turn_id: turn_context.sub_id.clone(),
trace_id: turn_context.trace_id.clone(),
started_at: turn_context.turn_timing_state.started_at_unix_secs().await,
model_context_window: turn_context.model_context_window(),
collaboration_mode_kind: turn_context.collaboration_mode.mode,
});
sess.send_event(&turn_context, start_event).await;
let world_state = Arc::new(
sess.build_world_state_for_environments(&turn_context, &turn_context.environments)
.await,
);
run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Manual).await
}
/// Runs token-budget inline auto-compaction as a normal compaction lifecycle.
///
/// Token-budget compaction skips model/server summarization and installs a fresh context window
/// instead. It is still modeled as compaction so compact hooks and `ContextCompaction` turn items
/// observe the same lifecycle as local or remote compaction.
pub(crate) async fn run_inline_auto_compact_task(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
world_state: Arc<WorldState>,
) -> CodexResult<()> {
run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Auto).await
}
async fn run_compact_task_inner(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
world_state: Arc<WorldState>,
trigger: CompactionTrigger,
) -> CodexResult<()> {
let pre_compact_outcome = run_pre_compact_hooks(sess, turn_context, trigger).await;
match pre_compact_outcome {
PreCompactHookOutcome::Continue => {}
PreCompactHookOutcome::Stopped => return Err(CodexErr::TurnAborted),
}
let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new());
sess.emit_turn_item_started(turn_context, &compaction_item)
.await;
sess.start_new_context_window(
turn_context.as_ref(),
world_state,
NewContextWindowMode::ForceStart,
)
.await;
sess.emit_turn_item_completed(turn_context, compaction_item)
.await;
let post_compact_outcome = run_post_compact_hooks(sess, turn_context, trigger).await;
if let PostCompactHookOutcome::Stopped = post_compact_outcome {
return Err(CodexErr::TurnAborted);
}
Ok(())
}
codex-rs/core/src/lib.rs
+1
View File
@@ -21,6 +21,7 @@ pub use turn_metadata::detached_memory_responses_metadata;
mod codex_thread;
mod compact_remote;
mod compact_remote_v2;
mod compact_token_budget;
mod config_lock;
pub use codex_thread::BackgroundTerminalInfo;
pub use codex_thread::CodexThread;
codex-rs/core/src/session/mod.rs
+29 -3
View File
@@ -253,6 +253,14 @@ use self::world_state::build_world_state_from_turn_context_item;
#[cfg(test)]
mod rollout_reconstruction_tests;
#[derive(Clone, Copy, Debug)]
pub(crate) enum NewContextWindowMode {
/// Start a fresh context window regardless of whether a new-context request is pending.
ForceStart,
/// Start a fresh context window only if one was explicitly requested.
StartIfRequested,
}
#[derive(Debug, PartialEq)]
pub enum SteerInputError {
NoActiveTurn(Vec<UserInput>),
@@ -3424,25 +3432,44 @@ impl Session {
state.request_new_context_window();
}
pub(crate) async fn maybe_start_new_context_window(
pub(crate) async fn start_new_context_window(
&self,
turn_context: &TurnContext,
world_state: Arc<WorldState>,
mode: NewContextWindowMode,
) -> Option<u64> {
let window = {
let mut state = self.state.lock().await;
state.start_new_context_window_if_requested()
state.start_new_context_window(mode)
};
let (window_number, window_ids) = window?;
self.install_new_context_window(turn_context, world_state, window_number, window_ids)
.await;
Some(window_number)
}
async fn install_new_context_window(
&self,
turn_context: &TurnContext,
world_state: Arc<WorldState>,
window_number: u64,
window_ids: AutoCompactWindowIds,
) {
let context_items = self
.build_initial_context_with_world_state(turn_context, world_state.as_ref())
.await;
let turn_context_item = turn_context.to_turn_context_item();
let previous_turn_settings = PreviousTurnSettings {
model: turn_context.model_info.slug.clone(),
comp_hash: turn_context.model_info.comp_hash.clone(),
realtime_active: Some(turn_context.realtime_active),
};
let replacement_history = context_items;
{
let mut state = self.state.lock().await;
state.replace_history(replacement_history.clone(), Some(turn_context_item.clone()));
state.history.set_world_state_baseline(world_state);
state.set_previous_turn_settings(Some(previous_turn_settings));
};
self.persist_rollout_items(&[
RolloutItem::Compacted(CompactedItem {
@@ -3461,7 +3488,6 @@ impl Session {
state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact);
}
self.recompute_token_usage(turn_context).await;
Some(window_number)
}
pub(crate) async fn reference_context_item(&self) -> Option<TurnContextItem> {
codex-rs/core/src/session/rollout_reconstruction.rs
+23 -8
View File
@@ -44,6 +44,7 @@ enum TurnReferenceContextItem {
struct ActiveReplaySegment<'a> {
turn_id: Option<String>,
counts_as_user_turn: bool,
has_compaction: bool,
previous_turn_settings: Option<PreviousTurnSettings>,
reference_context_item: TurnReferenceContextItem,
base_replacement_history: Option<&'a [ResponseItem]>,
@@ -85,19 +86,32 @@ fn finalize_active_segment<'a>(
*window = active_segment.window;
}
// `previous_turn_settings` come from the newest surviving user turn that established them.
if previous_turn_settings.is_none() && active_segment.counts_as_user_turn {
// `previous_turn_settings` come from the newest surviving user turn, or a compaction segment
// that established a fresh baseline with a `TurnContextItem`.
if previous_turn_settings.is_none()
&& (active_segment.counts_as_user_turn
|| (active_segment.has_compaction
&& matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Latest(_)
)))
{
*previous_turn_settings = active_segment.previous_turn_settings;
}
// `reference_context_item` comes from the newest surviving user turn baseline, or
// from a surviving compaction that explicitly cleared that baseline.
// `reference_context_item` comes from the newest surviving user turn baseline, from a
// surviving compaction that explicitly cleared that baseline, or from a compaction segment
// that re-established the baseline with a fresh `TurnContextItem`.
if matches!(reference_context_item, TurnReferenceContextItem::NeverSet)
&& (active_segment.counts_as_user_turn
|| matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Cleared
))
|| (active_segment.has_compaction
&& (matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Latest(_)
) || matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Cleared
))))
{
*reference_context_item = active_segment.reference_context_item;
}
@@ -149,6 +163,7 @@ impl Session {
RolloutItem::Compacted(compacted) => {
let active_segment =
active_segment.get_or_insert_with(ActiveReplaySegment::default);
active_segment.has_compaction = true;
if active_segment.window.is_none()
&& let Some(window_number) = compacted.window_number
{
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+40
View File
@@ -127,6 +127,46 @@ async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previ
assert!(session.reference_context_item().await.is_none());
}
#[tokio::test]
async fn reconstruct_history_hydrates_reference_context_from_compaction_turn_context() {
let (session, turn_context) = make_session_and_context().await;
let mut compact_context_item = turn_context.to_turn_context_item();
compact_context_item.turn_id = Some("compact-turn".to_string());
compact_context_item.model = "compact-model".to_string();
let replacement_history = vec![assistant_message("fresh context baseline")];
let rollout_items = vec![
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(replacement_history.clone()),
window_number: Some(1),
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
RolloutItem::TurnContext(compact_context_item.clone()),
];
let reconstructed = session
.reconstruct_history_from_rollout(&turn_context, &rollout_items)
.await;
assert_eq!(reconstructed.history, replacement_history);
assert_eq!(
reconstructed.previous_turn_settings,
Some(PreviousTurnSettings {
model: "compact-model".to_string(),
comp_hash: None,
realtime_active: Some(turn_context.realtime_active),
})
);
assert_eq!(
serde_json::to_value(reconstructed.reference_context_item)
.expect("serialize reconstructed reference context item"),
serde_json::to_value(Some(compact_context_item))
.expect("serialize expected reference context item")
);
}
#[tokio::test]
async fn record_initial_history_resumed_hydrates_previous_turn_settings_from_lifecycle_turn_with_missing_turn_context_id()
{
codex-rs/core/src/session/turn.rs
+107 -28
View File
@@ -39,6 +39,7 @@ use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::CodexResponsesRequestKind;
use crate::responses_retry::ResponsesStreamRequest;
use crate::responses_retry::handle_retryable_response_stream_error;
use crate::session::NewContextWindowMode;
use crate::session::PreviousTurnSettings;
use crate::session::TurnInput;
use crate::session::session::Session;
@@ -125,6 +126,20 @@ use tracing::trace;
use tracing::trace_span;
use tracing::warn;
#[derive(Clone, Copy, Debug)]
enum AutoCompactFollowUpState {
ActiveModelFollowUp,
NoActiveModelFollowUp,
}
#[derive(Debug)]
struct AutoCompactRequest {
initial_context_injection: InitialContextInjection,
reason: CompactionReason,
phase: CompactionPhase,
follow_up_state: AutoCompactFollowUpState,
}
/// Takes initial turn input and runs a loop where, at each sampling request,
/// the model replies with either:
///
@@ -340,7 +355,11 @@ pub(crate) async fn run_turn(
.await;
let started_new_context_window = sess
.maybe_start_new_context_window(turn_context.as_ref(), Arc::clone(&world_state))
.start_new_context_window(
turn_context.as_ref(),
Arc::clone(&world_state),
NewContextWindowMode::StartIfRequested,
)
.await
.is_some();
if started_new_context_window && needs_follow_up {
@@ -356,13 +375,24 @@ pub(crate) async fn run_turn(
&& token_limit_reached
&& needs_follow_up
{
let follow_up_state = if model_needs_follow_up {
AutoCompactFollowUpState::ActiveModelFollowUp
} else {
AutoCompactFollowUpState::NoActiveModelFollowUp
};
if let Err(err) = run_auto_compact(
&sess,
Arc::clone(&step_context),
&mut client_session,
InitialContextInjection::BeforeLastUserMessage(Arc::clone(&world_state)),
CompactionReason::ContextLimit,
CompactionPhase::MidTurn,
AutoCompactRequest {
initial_context_injection:
InitialContextInjection::BeforeLastUserMessage(Arc::clone(
&world_state,
)),
reason: CompactionReason::ContextLimit,
phase: CompactionPhase::MidTurn,
follow_up_state,
},
)
.await
{
@@ -886,9 +916,12 @@ async fn run_pre_sampling_compact(
sess,
step_context,
client_session,
InitialContextInjection::DoNotInject,
CompactionReason::ContextLimit,
CompactionPhase::PreTurn,
AutoCompactRequest {
initial_context_injection: InitialContextInjection::DoNotInject,
reason: CompactionReason::ContextLimit,
phase: CompactionPhase::PreTurn,
follow_up_state: AutoCompactFollowUpState::NoActiveModelFollowUp,
},
)
.await?;
}
@@ -926,16 +959,24 @@ async fn maybe_run_previous_model_inline_compact(
);
if should_compact_for_comp_hash_change {
let compact_turn_context = if turn_context.config.features.enabled(Feature::TokenBudget) {
turn_context
} else {
&previous_model_turn_context
};
let step_context = sess
.capture_step_context(Arc::clone(&previous_model_turn_context))
.capture_step_context(Arc::clone(compact_turn_context))
.await;
run_auto_compact(
sess,
step_context,
client_session,
InitialContextInjection::DoNotInject,
CompactionReason::CompHashChanged,
CompactionPhase::PreTurn,
AutoCompactRequest {
initial_context_injection: InitialContextInjection::DoNotInject,
reason: CompactionReason::CompHashChanged,
phase: CompactionPhase::PreTurn,
follow_up_state: AutoCompactFollowUpState::NoActiveModelFollowUp,
},
)
.await?;
return Ok(());
@@ -966,16 +1007,24 @@ async fn maybe_run_previous_model_inline_compact(
&& previous_model_turn_context.model_info.slug != turn_context.model_info.slug
&& old_context_window > new_context_window;
if should_run {
let compact_turn_context = if turn_context.config.features.enabled(Feature::TokenBudget) {
turn_context
} else {
&previous_model_turn_context
};
let step_context = sess
.capture_step_context(Arc::clone(&previous_model_turn_context))
.capture_step_context(Arc::clone(compact_turn_context))
.await;
run_auto_compact(
sess,
step_context,
client_session,
InitialContextInjection::DoNotInject,
CompactionReason::ModelDownshift,
CompactionPhase::PreTurn,
AutoCompactRequest {
initial_context_injection: InitialContextInjection::DoNotInject,
reason: CompactionReason::ModelDownshift,
phase: CompactionPhase::PreTurn,
follow_up_state: AutoCompactFollowUpState::NoActiveModelFollowUp,
},
)
.await?;
}
@@ -985,17 +1034,47 @@ async fn maybe_run_previous_model_inline_compact(
#[instrument(
level = "trace",
skip_all,
fields(reason = ?reason, phase = ?phase)
fields(reason = ?request.reason, phase = ?request.phase)
)]
async fn run_auto_compact(
sess: &Arc<Session>,
step_context: Arc<StepContext>,
client_session: &mut ModelClientSession,
initial_context_injection: InitialContextInjection,
reason: CompactionReason,
phase: CompactionPhase,
request: AutoCompactRequest,
) -> CodexResult<()> {
let turn_context = &step_context.turn;
if turn_context.config.features.enabled(Feature::TokenBudget) {
if matches!(request.phase, CompactionPhase::MidTurn)
&& matches!(
request.follow_up_state,
AutoCompactFollowUpState::ActiveModelFollowUp
)
{
// A token-budget reset clears the current-window history. Preserve active tool-call
// continuation state first; pre-turn compaction will reset the window before the next
// user turn if the budget is still exhausted.
return Ok(());
}
// Compaction is the reset request, so force a new context window
// instead of consuming a pending `new_context` tool request.
let world_state = match &request.initial_context_injection {
InitialContextInjection::BeforeLastUserMessage(world_state) => Arc::clone(world_state),
InitialContextInjection::DoNotInject => Arc::new(
sess.build_world_state_for_environments(
turn_context.as_ref(),
&step_context.environments,
)
.await,
),
};
crate::compact_token_budget::run_inline_auto_compact_task(
Arc::clone(sess),
Arc::clone(turn_context),
world_state,
)
.await?;
return Ok(());
}
if should_use_remote_compact_task(turn_context.provider.info()) {
if turn_context
.config
@@ -1011,9 +1090,9 @@ async fn run_auto_compact(
Arc::clone(sess),
step_context,
client_session,
initial_context_injection,
reason,
phase,
request.initial_context_injection,
request.reason,
request.phase,
)
.await?;
return Ok(());
@@ -1027,9 +1106,9 @@ async fn run_auto_compact(
Arc::clone(sess),
step_context,
client_session.turn_state(),
initial_context_injection,
reason,
phase,
request.initial_context_injection,
request.reason,
request.phase,
)
.await?;
} else {
@@ -1041,9 +1120,9 @@ async fn run_auto_compact(
run_inline_auto_compact_task(
Arc::clone(sess),
Arc::clone(turn_context),
initial_context_injection,
reason,
phase,
request.initial_context_injection,
request.reason,
request.phase,
)
.await?;
}
codex-rs/core/src/state/session.rs
+6 -2
View File
@@ -12,6 +12,7 @@ use super::auto_compact_window::AutoCompactWindow;
use super::auto_compact_window::AutoCompactWindowIds;
use super::auto_compact_window::AutoCompactWindowSnapshot;
use crate::context_manager::ContextManager;
use crate::session::NewContextWindowMode;
use crate::session::PreviousTurnSettings;
use crate::session::session::SessionConfiguration;
use crate::session::time_reminder::CurrentTimeReminderState;
@@ -187,10 +188,13 @@ impl SessionState {
self.auto_compact_window.request_new_context_window();
}
pub(crate) fn start_new_context_window_if_requested(
pub(crate) fn start_new_context_window(
&mut self,
mode: NewContextWindowMode,
) -> Option<(u64, AutoCompactWindowIds)> {
if !self.auto_compact_window.take_new_context_window_request() {
if matches!(mode, NewContextWindowMode::StartIfRequested)
&& !self.auto_compact_window.take_new_context_window_request()
{
return None;
}
codex-rs/core/src/tasks/compact.rs
+6
View File
@@ -7,6 +7,7 @@ use super::emit_compact_metric;
use crate::session::TurnInput;
use crate::session::turn_context::TurnContext;
use crate::state::TaskKind;
use codex_features::Feature;
use codex_protocol::error::CodexErr;
use codex_protocol::user_input::UserInput;
use tokio_util::sync::CancellationToken;
@@ -31,6 +32,11 @@ impl SessionTask for CompactTask {
_cancellation_token: CancellationToken,
) -> SessionTaskResult {
let session = session.clone_session();
if ctx.config.features.enabled(Feature::TokenBudget) {
crate::compact_token_budget::run_manual_compact_task(session, ctx).await?;
return Ok(None);
}
let result = if crate::compact::should_use_remote_compact_task(ctx.provider.info()) {
if ctx
.config
codex-rs/core/tests/suite/token_budget.rs
+531 -11
View File
@@ -3,34 +3,56 @@ use codex_config::types::McpServerConfig;
use codex_config::types::McpServerTransportConfig;
use codex_core::config::TokenBudgetConfig;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_model_provider_info::built_in_model_providers;
use codex_models_manager::bundled_models_response;
use codex_protocol::config_types::CollaborationMode;
use codex_protocol::config_types::ModeKind;
use codex_protocol::config_types::Settings;
use codex_protocol::items::TurnItem;
use codex_protocol::models::PermissionProfile;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::CONTEXT_WINDOW_CLOSE_TAG;
use codex_protocol::protocol::CONTEXT_WINDOW_OPEN_TAG;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::HookEventName;
use codex_protocol::protocol::HookRunStatus;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::ItemStartedEvent;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::ThreadSettingsOverrides;
use core_test_support::PathBufExt;
use core_test_support::assert_regex_match;
use core_test_support::context_snapshot;
use core_test_support::context_snapshot::ContextSnapshotOptions;
use core_test_support::hooks::trust_discovered_hooks;
use core_test_support::responses::ResponsesRequest;
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_compact_json_once;
use core_test_support::responses::mount_sse_sequence;
use core_test_support::responses::sse;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::stdio_server_bin;
use core_test_support::test_codex::local;
use core_test_support::test_codex::local_selections;
use core_test_support::test_codex::test_codex;
use core_test_support::test_codex::turn_permission_fields;
use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use core_test_support::wait_for_mcp_server;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;
use std::time::Duration;
const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
@@ -79,6 +101,148 @@ fn tool_names(request: &ResponsesRequest) -> Vec<String> {
.collect()
}
fn model_info_with_comp_hash(slug: &str, comp_hash: &str) -> ModelInfo {
let mut model_info = bundled_models_response()
.expect("bundled models.json should parse")
.models
.into_iter()
.find(|model_info| model_info.slug == slug)
.expect("model missing from models.json");
model_info.comp_hash = Some(comp_hash.to_string());
model_info
}
fn disabled_permission_user_turn(text: impl Into<String>, cwd: PathBuf, model: String) -> Op {
let (sandbox_policy, permission_profile) =
turn_permission_fields(PermissionProfile::Disabled, cwd.as_path());
Op::UserInput {
items: vec![codex_protocol::user_input::UserInput::Text {
text: text.into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: ThreadSettingsOverrides {
environments: Some(local_selections(cwd.abs())),
approval_policy: Some(AskForApproval::Never),
sandbox_policy: Some(sandbox_policy),
permission_profile,
collaboration_mode: Some(CollaborationMode {
mode: ModeKind::Default,
settings: Settings {
model,
reasoning_effort: None,
developer_instructions: None,
},
}),
..Default::default()
},
}
}
fn python_hook_command(script_path: &Path) -> String {
format!("python3 \"{}\"", script_path.display())
}
fn write_token_budget_compact_hooks(home: &Path) {
let script_path = home.join("token_budget_compact_hook.py");
let log_path = home.join("token_budget_compact_hook_log.jsonl");
let script = format!(
r#"import json
from pathlib import Path
import sys
payload = json.load(sys.stdin)
with Path(r"{log_path}").open("a", encoding="utf-8") as handle:
handle.write(json.dumps(payload) + "\n")
"#,
log_path = log_path.display(),
);
std::fs::write(&script_path, script).expect("write compact hook script");
let hooks = json!({
"hooks": {
"PreCompact": [{
"matcher": "manual",
"hooks": [{
"type": "command",
"command": python_hook_command(&script_path),
}]
}, {
"matcher": "auto",
"hooks": [{
"type": "command",
"command": python_hook_command(&script_path),
}]
}],
"PostCompact": [{
"matcher": "manual",
"hooks": [{
"type": "command",
"command": python_hook_command(&script_path),
}]
}, {
"matcher": "auto",
"hooks": [{
"type": "command",
"command": python_hook_command(&script_path),
}]
}]
}
});
std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json");
}
fn read_hook_inputs(path: &Path) -> Vec<Value> {
let contents = std::fs::read_to_string(path).expect("read hook log");
contents
.lines()
.map(|line| serde_json::from_str(line).expect("parse hook payload"))
.collect()
}
async fn assert_context_compaction_item_lifecycle(codex: &std::sync::Arc<codex_core::CodexThread>) {
let mut started_id = None;
let mut completed_id = None;
loop {
let event = codex.next_event().await.expect("next event");
match event.msg {
EventMsg::ItemStarted(ItemStartedEvent {
item: TurnItem::ContextCompaction(item),
..
}) => {
assert!(started_id.is_none(), "compaction item should start once");
assert!(
completed_id.is_none(),
"compaction item should not complete before it starts"
);
started_id = Some(item.id);
}
EventMsg::ItemCompleted(ItemCompletedEvent {
item: TurnItem::ContextCompaction(item),
..
}) => {
assert!(
completed_id.is_none(),
"compaction item should complete once"
);
assert_eq!(
started_id.as_deref(),
Some(item.id.as_str()),
"compaction item completion should match the started item"
);
completed_id = Some(item.id);
}
EventMsg::TurnComplete(_) => break,
_ => {}
}
}
assert!(started_id.is_some(), "compaction item should start");
assert!(completed_id.is_some(), "compaction item should complete");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> {
skip_if_no_network!(Ok(()));
@@ -395,19 +559,18 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
sse(vec![
ev_response_created("resp-compact"),
ev_assistant_message("msg-compact", "compact summary"),
ev_completed("resp-compact"),
ev_response_created("resp-1"),
ev_assistant_message("msg-1", "assistant before compact"),
ev_completed("resp-1"),
]),
sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let compact = mount_compact_json_once(&server, json!({ "output": [] })).await;
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.name = "OpenAI-compatible test provider".to_string();
model_provider.base_url = Some(format!("{}/v1", server.uri()));
model_provider.supports_websockets = false;
@@ -425,21 +588,22 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
test.submit_turn("before compact").await?;
test.codex.submit(Op::Compact).await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
assert_context_compaction_item_lifecycle(&test.codex).await;
test.submit_turn("after compact").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 3);
assert_eq!(requests.len(), 2);
assert!(
compact.requests().is_empty(),
"token budget compaction should not call server-side compaction"
);
let thread_id = test.session_configured.thread_id;
let initial_token_budget = token_budget_contexts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let (initial_first_window_id, initial_previous_window_id, initial_window_id) =
token_budget_window_ids(&initial_token_budget[0], thread_id);
let post_compaction_token_budget = token_budget_contexts(&requests[2]);
let post_compaction_token_budget = token_budget_contexts(&requests[1]);
assert_eq!(post_compaction_token_budget.len(), 1);
let (
post_compaction_first_window_id,
@@ -454,6 +618,362 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
Some(initial_window_id.as_str())
);
assert_ne!(post_compaction_window_id, initial_window_id);
assert!(
!requests[1].body_contains_text("before compact"),
"token budget compaction should drop prior user messages"
);
assert!(
!requests[1].body_contains_text("assistant before compact"),
"token budget compaction should drop prior assistant messages"
);
assert!(
requests[1].body_contains_text("after compact"),
"follow-up should still include the new turn input"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_compaction_runs_compact_hooks() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let test = test_codex()
.with_pre_build_hook(write_token_budget_compact_hooks)
.with_config(|config| {
config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
trust_discovered_hooks(config);
})
.build(&server)
.await?;
test.codex.submit(Op::Compact).await?;
let pre_compact = wait_for_event_match(&test.codex, |event| match event {
EventMsg::HookCompleted(completed)
if completed.run.event_name == HookEventName::PreCompact =>
{
Some(completed.clone())
}
_ => None,
})
.await;
assert_eq!(pre_compact.run.status, HookRunStatus::Completed);
let post_compact = wait_for_event_match(&test.codex, |event| match event {
EventMsg::HookCompleted(completed)
if completed.run.event_name == HookEventName::PostCompact =>
{
Some(completed.clone())
}
_ => None,
})
.await;
assert_eq!(post_compact.run.status, HookRunStatus::Completed);
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let hook_inputs = read_hook_inputs(
&test
.codex_home_path()
.join("token_budget_compact_hook_log.jsonl"),
);
assert_eq!(hook_inputs.len(), 2);
assert_eq!(hook_inputs[0]["hook_event_name"], "PreCompact");
assert_eq!(hook_inputs[0]["trigger"], "manual");
assert_eq!(hook_inputs[1]["hook_event_name"], "PostCompact");
assert_eq!(hook_inputs[1]["trigger"], "manual");
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_pre_turn_auto_compaction_starts_new_window() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_assistant_message("msg-1", "assistant before auto compact"),
ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500),
]),
sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let compact = mount_compact_json_once(&server, json!({ "output": [] })).await;
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.base_url = Some(format!("{}/v1", server.uri()));
model_provider.supports_websockets = false;
let test = test_codex()
.with_pre_build_hook(write_token_budget_compact_hooks)
.with_config(move |config| {
config.model_provider = model_provider;
config.model_context_window = Some(10_000);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
trust_discovered_hooks(config);
})
.build(&server)
.await?;
test.submit_turn("before auto compact").await?;
test.submit_turn("after auto compact").await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert!(
compact.requests().is_empty(),
"token budget auto-compaction should not call server-side compaction"
);
let thread_id = test.session_configured.thread_id;
let initial_token_budget = token_budget_contexts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let (initial_first_window_id, initial_previous_window_id, initial_window_id) =
token_budget_window_ids(&initial_token_budget[0], thread_id);
let post_compaction_token_budget = token_budget_contexts(&requests[1]);
assert_eq!(post_compaction_token_budget.len(), 1);
let (
post_compaction_first_window_id,
post_compaction_previous_window_id,
post_compaction_window_id,
) = token_budget_window_ids(&post_compaction_token_budget[0], thread_id);
assert_eq!(initial_previous_window_id, None);
assert_eq!(initial_first_window_id, initial_window_id);
assert_eq!(post_compaction_first_window_id, initial_first_window_id);
assert_eq!(
post_compaction_previous_window_id.as_deref(),
Some(initial_window_id.as_str())
);
assert_ne!(post_compaction_window_id, initial_window_id);
assert!(
!requests[1].body_contains_text("before auto compact"),
"token budget auto-compaction should drop prior user messages"
);
assert!(
!requests[1].body_contains_text("assistant before auto compact"),
"token budget auto-compaction should drop prior assistant messages"
);
assert!(
requests[1].body_contains_text("after auto compact"),
"follow-up should still include the new turn input"
);
let hook_inputs = read_hook_inputs(
&test
.codex_home_path()
.join("token_budget_compact_hook_log.jsonl"),
);
assert_eq!(hook_inputs.len(), 2);
assert_eq!(hook_inputs[0]["hook_event_name"], "PreCompact");
assert_eq!(hook_inputs[0]["trigger"], "auto");
assert_eq!(hook_inputs[1]["hook_event_name"], "PostCompact");
assert_eq!(hook_inputs[1]["trigger"], "auto");
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_comp_hash_change_uses_current_model_context_for_new_window() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let previous_model = "gpt-5.3-codex";
let next_model = "gpt-5.2";
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_assistant_message("msg-1", "before switch"),
ev_completed_with_tokens("resp-1", /*total_tokens*/ 100),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-2", "after switch"),
ev_completed("resp-2"),
]),
],
)
.await;
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.name = "OpenAI (test)".into();
model_provider.base_url = Some(format!("{}/v1", server.uri()));
model_provider.supports_websockets = false;
let test = test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_model(previous_model)
.with_config(move |config| {
config.model_provider = model_provider;
config.model_catalog = Some(ModelsResponse {
models: vec![
model_info_with_comp_hash(previous_model, "hash-a"),
model_info_with_comp_hash(next_model, "hash-b"),
],
});
config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.codex
.submit(disabled_permission_user_turn(
"before switch",
test.cwd.path().to_path_buf(),
previous_model.to_string(),
))
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
test.codex
.submit(disabled_permission_user_turn(
"after switch",
test.cwd.path().to_path_buf(),
next_model.to_string(),
))
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let requests = responses.requests();
assert_eq!(
requests.len(),
2,
"token-budget comp-hash changes should reset locally before sampling the next turn"
);
assert_eq!(
requests[0].body_json()["model"].as_str(),
Some(previous_model)
);
assert_eq!(requests[1].body_json()["model"].as_str(), Some(next_model));
let thread_id = test.session_configured.thread_id;
let initial_token_budget = token_budget_contexts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let (initial_first_window_id, _, initial_window_id) =
token_budget_window_ids(&initial_token_budget[0], thread_id);
let post_compaction_token_budget = token_budget_contexts(&requests[1]);
assert_eq!(post_compaction_token_budget.len(), 1);
let (
post_compaction_first_window_id,
post_compaction_previous_window_id,
post_compaction_window_id,
) = token_budget_window_ids(&post_compaction_token_budget[0], thread_id);
assert_eq!(post_compaction_first_window_id, initial_first_window_id);
assert_eq!(
post_compaction_previous_window_id.as_deref(),
Some(initial_window_id.as_str())
);
assert_ne!(post_compaction_window_id, initial_window_id);
assert!(
!requests[1].body_contains_text("before switch"),
"token-budget comp-hash changes should drop prior user messages"
);
assert_eq!(
requests[1]
.message_input_texts("developer")
.into_iter()
.filter(|text| text.contains("<model_switch>"))
.count(),
1,
"fresh token-budget window should be built from the current turn context without duplicating the model switch"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_mid_turn_auto_compaction_preserves_active_follow_up() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "remaining-call";
let responses = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_function_call(call_id, "get_context_remaining", "{}"),
ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-2", "done"),
ev_completed("resp-2"),
]),
],
)
.await;
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.name = "OpenAI (test)".into();
model_provider.base_url = Some(format!("{}/v1", server.uri()));
model_provider.supports_websockets = false;
let test = test_codex()
.with_config(move |config| {
config.model_provider = model_provider;
config.model_context_window = Some(10_000);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
test.submit_turn("trigger mid-turn auto compaction").await?;
let requests = responses.requests();
assert_eq!(
requests.len(),
2,
"active model follow-up should not run a compaction request before the continuation"
);
assert_eq!(
requests[1].function_call_output_content_and_success(call_id),
Some((
Some("You have 0 tokens left in this context window.".to_string()),
None
)),
"follow-up request should preserve active tool output after mid-turn compaction"
);
let thread_id = test.session_configured.thread_id;
let initial_token_budget = token_budget_contexts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let (initial_first_window_id, _, initial_window_id) =
token_budget_window_ids(&initial_token_budget[0], thread_id);
let follow_up_token_budget = token_budget_contexts(&requests[1]);
assert_eq!(follow_up_token_budget.len(), 1);
let (follow_up_first_window_id, follow_up_previous_window_id, follow_up_window_id) =
token_budget_window_ids(&follow_up_token_budget[0], thread_id);
assert_eq!(follow_up_first_window_id, initial_first_window_id);
assert_eq!(follow_up_previous_window_id, None);
assert_eq!(
follow_up_window_id, initial_window_id,
"mid-turn auto-compaction should not reset the token-budget context window while the model needs follow-up"
);
Ok(())
}