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 000000000..0ff7506bc --- /dev/null +++ b/codex-rs/core/src/compact_token_budget.rs @@ -0,0 +1,90 @@ +use std::sync::Arc; + +use crate::compact::InitialContextInjection; +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::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, + initial_context_injection: InitialContextInjection, +) -> CodexResult<()> { + let world_state = match initial_context_injection { + InitialContextInjection::BeforeLastUserMessage(world_state) => world_state, + InitialContextInjection::DoNotInject => Arc::new( + sess.build_world_state_for_environments(&turn_context, &turn_context.environments) + .await, + ), + }; + 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) + .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 121c8d2be..519d11d1c 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 63ebf6729..d94ca7f34 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3425,16 +3425,21 @@ impl Session { state.request_new_context_window(); } - pub(crate) async fn maybe_start_new_context_window( + pub(crate) async fn take_new_context_window_request(&self) -> bool { + let mut state = self.state.lock().await; + state.take_new_context_window_request() + } + + pub(crate) async fn start_new_context_window( &self, turn_context: &TurnContext, world_state: Arc, - ) -> Option { + ) -> u64 { let window = { let mut state = self.state.lock().await; - state.start_new_context_window_if_requested() + state.start_new_context_window() }; - let (window_number, window_ids) = window?; + let (window_number, window_ids) = window; let context_items = self .build_initial_context_with_world_state(turn_context, world_state.as_ref()) .await; @@ -3462,7 +3467,7 @@ impl Session { state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); } self.recompute_token_usage(turn_context).await; - Some(window_number) + window_number } pub(crate) async fn reference_context_item(&self) -> Option { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2ca716817..397dd17b6 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -331,22 +331,14 @@ 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)) - .await - .is_some(); - if started_new_context_window && needs_follow_up { - can_drain_pending_input = !model_needs_follow_up; - continue; - } - // 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 turn_context + let auto_compact_needed = turn_context .config .features .enabled(Feature::AutoCompaction) - && token_limit_reached - && needs_follow_up + && token_limit_reached; + if needs_follow_up + && (sess.take_new_context_window_request().await || auto_compact_needed) { if let Err(err) = run_auto_compact( &sess, @@ -928,6 +920,18 @@ async fn run_auto_compact( phase: CompactionPhase, ) -> CodexResult<()> { let turn_context = &step_context.turn; + if turn_context.config.features.enabled(Feature::TokenBudget) { + // Compaction is the reset request, so force a new context window + // instead of consuming a pending `new_context` tool request. + crate::compact_token_budget::run_inline_auto_compact_task( + Arc::clone(sess), + Arc::clone(turn_context), + initial_context_injection, + ) + .await?; + return Ok(()); + } + if should_use_remote_compact_task(turn_context.provider.info()) { if turn_context .config diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 521559486..d681c6fa8 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -187,16 +187,14 @@ impl SessionState { self.auto_compact_window.request_new_context_window(); } - pub(crate) fn start_new_context_window_if_requested( - &mut self, - ) -> Option<(u64, AutoCompactWindowIds)> { - if !self.auto_compact_window.take_new_context_window_request() { - return None; - } + pub(crate) fn take_new_context_window_request(&mut self) -> bool { + self.auto_compact_window.take_new_context_window_request() + } + pub(crate) fn start_new_context_window(&mut self) -> (u64, AutoCompactWindowIds) { let window = self.auto_compact_window.advance(); self.auto_compact_window.clear_prefill(); - Some(window) + window } pub(crate) fn token_info(&self) -> Option { diff --git a/codex-rs/core/src/tasks/compact.rs b/codex-rs/core/src/tasks/compact.rs index 2c45fd8ee..60cdebb82 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 319965086..a4aac9d93 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -5,20 +5,27 @@ use codex_core::config::TokenBudgetConfig; use codex_features::Feature; use codex_model_provider_info::built_in_model_providers; use codex_protocol::config_types::AutoCompactTokenLimitScope; +use codex_protocol::items::TurnItem; 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 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; @@ -27,11 +34,13 @@ use core_test_support::stdio_server_bin; use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; 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::time::Duration; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; @@ -96,6 +105,62 @@ fn ev_completed_with_usage(id: &str, input_tokens: i64, output_tokens: i64) -> V }) } +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"); + std::fs::write( + &script_path, + "import json\nimport sys\njson.load(sys.stdin)\n", + ) + .expect("write compact hook script"); + let hooks = json!({ + "hooks": { + "PreCompact": [{ + "matcher": "manual", + "hooks": [{ + "type": "command", + "command": python_hook_command(&script_path), + }] + }], + "PostCompact": [{ + "matcher": "manual", + "hooks": [{ + "type": "command", + "command": python_hook_command(&script_path), + }] + }] + } + }); + std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json"); +} + +async fn assert_context_compaction_item_lifecycle(codex: &std::sync::Arc) { + let mut saw_compaction_started = false; + let mut saw_compaction_completed = false; + + loop { + let event = codex.next_event().await.expect("next event"); + match event.msg { + EventMsg::ItemStarted(ItemStartedEvent { + item: TurnItem::ContextCompaction(_), + .. + }) => saw_compaction_started = true, + EventMsg::ItemCompleted(ItemCompletedEvent { + item: TurnItem::ContextCompaction(_), + .. + }) => saw_compaction_completed = true, + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + + assert!(saw_compaction_started, "compaction item should start"); + assert!(saw_compaction_completed, "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(())); @@ -540,19 +605,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; @@ -570,21 +634,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, @@ -599,6 +664,147 @@ 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; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_mid_turn_auto_compaction_resets_before_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, + "token-budget auto-compaction should reset locally before the continuation" + ); + assert!( + !requests[1].input().iter().any(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + && item.get("call_id").and_then(Value::as_str) == Some(call_id) + }), + "fresh token-budget windows should drop active tool output with the prior history" + ); + + 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.as_deref(), + Some(initial_window_id.as_str()) + ); + assert!( + !requests[1].body_contains_text("trigger mid-turn auto compaction"), + "fresh token-budget windows should drop prior user messages" + ); + assert_ne!( + follow_up_window_id, initial_window_id, + "mid-turn token-budget auto-compaction should reset the context window" + ); Ok(()) }