openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

core: fix token-budget compaction state

Browse Source
This commit is contained in:
Michael Bolin
2026-06-23 17:59:10 -07:00
parent 97dce078c5
commit 129472859a
4 changed files with 236 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
codex-rs/core/src/session/rollout_reconstruction.rs
View File

@@ -85,15 +85,26 @@ 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 {
let compact_segment_reference_context_item = active_segment.base_replacement_history.is_some()
&& matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Latest(_)
);
// `previous_turn_settings` come from the newest surviving user turn that established them, or
// from a compact replacement history that persisted the turn context used to build it.
if previous_turn_settings.is_none()
&& (active_segment.counts_as_user_turn || compact_segment_reference_context_item)
{
*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, a compact
// replacement history that persisted its new baseline, or a surviving compaction that
// explicitly cleared the previous baseline.
if matches!(reference_context_item, TurnReferenceContextItem::NeverSet)
&& (active_segment.counts_as_user_turn
|| compact_segment_reference_context_item
|| matches!(
active_segment.reference_context_item,
TurnReferenceContextItem::Cleared

60
codex-rs/core/src/session/rollout_reconstruction_tests.rs
View File

@@ -1111,6 +1111,66 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear
assert!(reconstructed.reference_context_item.is_none());
}
#[tokio::test]
async fn record_initial_history_resumed_standalone_compact_reestablishes_reference_context_item() {
let (session, turn_context) = make_session_and_context().await;
let compact_context_item = turn_context.to_turn_context_item();
let compact_turn_id = compact_context_item
.turn_id
.clone()
.expect("turn context should have turn_id");
let rollout_items = vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(
codex_protocol::protocol::TurnStartedEvent {
turn_id: compact_turn_id.clone(),
trace_id: None,
started_at: None,
model_context_window: Some(128_000),
collaboration_mode_kind: ModeKind::Default,
},
)),
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(vec![assistant_message("compacted baseline")]),
window_number: Some(1),
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
RolloutItem::TurnContext(compact_context_item.clone()),
RolloutItem::EventMsg(EventMsg::TurnComplete(
codex_protocol::protocol::TurnCompleteEvent {
turn_id: compact_turn_id,
last_agent_message: None,
completed_at: None,
duration_ms: None,
time_to_first_token_ms: None,
},
)),
];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: Arc::new(rollout_items),
rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")),
}))
.await;
assert_eq!(
session.previous_turn_settings().await,
Some(PreviousTurnSettings {
model: compact_context_item.model.clone(),
comp_hash: compact_context_item.comp_hash.clone(),
realtime_active: compact_context_item.realtime_active,
})
);
assert_eq!(
session.reference_context_item().await,
Some(compact_context_item)
);
}
#[tokio::test]
async fn record_initial_history_resumed_turn_context_after_compaction_reestablishes_reference_context_item()
{

9
codex-rs/core/src/session/turn.rs
View File

@@ -849,10 +849,15 @@ async fn maybe_run_previous_model_inline_compact(
.with_model(previous_turn_settings.model, &sess.services.models_manager)
.await,
);
let compaction_turn_context = if turn_context.config.features.enabled(Feature::TokenBudget) {
Arc::clone(turn_context)
} else {
Arc::clone(&previous_model_turn_context)
};
if should_compact_for_comp_hash_change {
let step_context = sess
.capture_step_context(Arc::clone(&previous_model_turn_context))
.capture_step_context(Arc::clone(&compaction_turn_context))
.await;
run_auto_compact(
sess,
@@ -892,7 +897,7 @@ async fn maybe_run_previous_model_inline_compact(
&& old_context_window > new_context_window;
if should_run {
let step_context = sess
.capture_step_context(Arc::clone(&previous_model_turn_context))
.capture_step_context(Arc::clone(&compaction_turn_context))
.await;
run_auto_compact(
sess,

154
codex-rs/core/tests/suite/token_budget.rs
View File

@@ -4,8 +4,13 @@ use codex_config::types::McpServerTransportConfig;
use codex_core::config::TokenBudgetConfig;
use codex_features::Feature;
use codex_model_provider_info::built_in_model_providers;
use codex_models_manager::bundled_models_response;
use codex_protocol::config_types::AutoCompactTokenLimitScope;
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;
@@ -14,6 +19,7 @@ use codex_protocol::protocol::HookRunStatus;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::ItemStartedEvent;
use codex_protocol::protocol::Op;
use codex_protocol::user_input::UserInput;
use core_test_support::PathBufExt;
use core_test_support::assert_regex_match;
use core_test_support::context_snapshot;
@@ -31,8 +37,11 @@ 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::TestCodex;
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;
@@ -41,10 +50,30 @@ use serde_json::Value;
use serde_json::json;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use tempfile::TempDir;
const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
fn model_info_with_context_window_and_comp_hash(
slug: &str,
context_window: i64,
comp_hash: &str,
) -> ModelInfo {
let models_response = bundled_models_response().expect("bundled models.json should parse");
let mut model_info = models_response
.models
.into_iter()
.find(|model| model.slug == slug)
.expect("model missing from models.json");
model_info.context_window = Some(context_window);
model_info.max_context_window = None;
model_info.comp_hash = Some(comp_hash.to_string());
model_info.effective_context_window_percent = 100;
model_info
}
fn token_budget_contexts(request: &ResponsesRequest) -> Vec<String> {
let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nThread id: ");
request
@@ -137,6 +166,42 @@ fn write_token_budget_compact_hooks(home: &Path) {
std::fs::write(home.join("hooks.json"), hooks.to_string()).expect("write hooks.json");
}
async fn submit_turn_with_model(test: &TestCodex, prompt: &str, model: &str) -> Result<()> {
let (sandbox_policy, permission_profile) =
turn_permission_fields(PermissionProfile::Disabled, test.config.cwd.as_path());
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: prompt.to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
environments: Some(local_selections(test.config.cwd.clone())),
approval_policy: Some(AskForApproval::Never),
sandbox_policy: Some(sandbox_policy),
permission_profile,
collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {
mode: codex_protocol::config_types::ModeKind::Default,
settings: codex_protocol::config_types::Settings {
model: model.to_string(),
reasoning_effort: None,
developer_instructions: None,
},
}),
..Default::default()
},
})
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
Ok(())
}
async fn assert_context_compaction_item_lifecycle(codex: &std::sync::Arc<codex_core::CodexThread>) {
let mut saw_compaction_started = false;
let mut saw_compaction_completed = false;
@@ -211,6 +276,95 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()>
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_comp_hash_reset_uses_current_model_context() -> 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_completed_with_tokens("resp-1", /*total_tokens*/ 100),
]),
sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let codex_home = Arc::new(TempDir::new()?);
let skill_dir = codex_home.path().join("skills/budget-sensitive");
std::fs::create_dir_all(&skill_dir)?;
std::fs::write(
skill_dir.join("SKILL.md"),
"---\nname: budget-sensitive-skill\ndescription: visible only with the current model budget\n---\n\n# Body\n",
)?;
let previous_model = "gpt-5.3-codex";
let current_model = "gpt-5.2";
let model_catalog = ModelsResponse {
models: vec![
model_info_with_context_window_and_comp_hash(
previous_model,
/*context_window*/ 100,
"hash-a",
),
model_info_with_context_window_and_comp_hash(
current_model,
/*context_window*/ 200_000,
"hash-b",
),
],
};
let codex_home_path = codex_home.path().to_path_buf();
let test = test_codex()
.with_home(codex_home)
.with_model(previous_model)
.with_config(move |config| {
config.cwd = codex_home_path.abs();
config.model_catalog = Some(model_catalog);
config
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
})
.build(&server)
.await?;
submit_turn_with_model(&test, "before comp hash switch", previous_model).await?;
submit_turn_with_model(&test, "after comp hash switch", current_model).await?;
let requests = responses.requests();
assert_eq!(requests.len(), 2);
assert_eq!(
requests[0].body_json()["model"].as_str(),
Some(previous_model)
);
assert_eq!(
requests[1].body_json()["model"].as_str(),
Some(current_model)
);
let initial_developer_text = requests[0].message_input_texts("developer").join("\n\n");
assert!(
!initial_developer_text.contains("budget-sensitive-skill"),
"previous model's tiny context budget should omit skill metadata"
);
let current_developer_text = requests[1].message_input_texts("developer").join("\n\n");
assert!(
current_developer_text
.contains("budget-sensitive-skill: visible only with the current model budget"),
"pre-turn token-budget reset should rebuild full context with the current model budget, got {current_developer_text:?}"
);
assert!(
!requests[1].body_contains_text("before comp hash switch"),
"token-budget comp-hash reset should drop prior window history"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> {
skip_if_no_network!(Ok(()));