diff --git a/codex-rs/core/src/compact_token_budget.rs b/codex-rs/core/src/compact_token_budget.rs new file mode 100644 index 0000000000..469e6d6fdc --- /dev/null +++ b/codex-rs/core/src/compact_token_budget.rs @@ -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, + turn_context: Arc, +) -> 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, + turn_context: Arc, + world_state: Arc, +) -> CodexResult<()> { + run_compact_task_inner(&sess, &turn_context, world_state, CompactionTrigger::Auto).await +} + +async fn run_compact_task_inner( + sess: &Arc, + turn_context: &Arc, + world_state: Arc, + 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(()) +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 121c8d2be2..519d11d1ca 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -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; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8b8044bad5..f4aaf0d4cd 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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), @@ -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, + mode: NewContextWindowMode, ) -> Option { 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, + 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 { diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index b3fbeb1038..3fc1675089 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -44,6 +44,7 @@ enum TurnReferenceContextItem { struct ActiveReplaySegment<'a> { turn_id: Option, counts_as_user_turn: bool, + has_compaction: bool, previous_turn_settings: Option, 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 { diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index b875e1f49e..a766270cf8 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -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() { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 100089723b..176f2497d3 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -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, step_context: Arc, 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?; } diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 5215594867..5182fcaf21 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -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; } diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index 2c45fd8ee2..60cdebb829 100644 --- a/codex-rs/core/src/tasks/compact.rs +++ b/codex-rs/core/src/tasks/compact.rs @@ -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 diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 09399e511e..6278cd9f7e 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -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 { .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, 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 { + 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) { + 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("")) + .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(()) }