From 2fa67f03f332a57b1bb34a8e14d55d14165fdcd3 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Mon, 2 Nov 2026 19:09:26 +0000 Subject: [PATCH] Allow volatile World State sections to refresh after compaction (#58956) ## Why Host-provided context such as budget reminders can become stale while compaction runs. Replaying a captured value or retaining its comparison snapshot can prevent the next step from emitting a fresh reminder, even when the remaining budget has changed. ## What changed - Add `WorldStateSectionContribution::without_full_context()` to omit a section's fragment and comparison snapshot from initial and compacted context rebuilds while preserving ordinary step updates. - Reset comparison state after prefix compaction for sections excluded at either capture or installation, including when their messages survive in the live suffix. - Include step World State updates when building debug prompt input. ## Testing Add regression coverage for budget reminders refreshing after remote compaction within the same threshold, unchanged request prefixes within a context window, refreshes with a retained live suffix, and step reminders in debug prompts. GitOrigin-RevId: ea4456755e16c93b5d16e8c762b26a33f8571a86 --- .../src/compact_remote_v2_installation.rs | 4 + codex-rs/core/src/context/world_state/mod.rs | 37 ++- codex-rs/core/src/prompt_debug.rs | 2 + codex-rs/core/src/session/mod.rs | 7 + .../core/tests/suite/prompt_debug_tests.rs | 75 ++++-- codex-rs/core/tests/suite/scenarios.rs | 42 +++- .../tests/suite/scenarios_budget_context.rs | 224 ++++++++++++++++++ ...tra_async_compaction_with_live_suffix.snap | 46 ++-- ..._host_budget_context_after_compaction.snap | 98 ++++++++ .../src/contributors/world_state.rs | 16 ++ 10 files changed, 503 insertions(+), 48 deletions(-) create mode 100644 codex-rs/core/tests/suite/scenarios_budget_context.rs create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap diff --git a/codex-rs/core/src/compact_remote_v2_installation.rs b/codex-rs/core/src/compact_remote_v2_installation.rs index a13b18ba44..8397c5d402 100644 --- a/codex-rs/core/src/compact_remote_v2_installation.rs +++ b/codex-rs/core/src/compact_remote_v2_installation.rs @@ -218,6 +218,10 @@ async fn build_and_install_compaction( prefix, prefix_context, initial_context, + reset_world_state_sections: world_state + .full_context_excluded_section_ids() + .chain(cutoff_world_state.full_context_excluded_section_ids()) + .collect(), model_info: &model_info, compaction_response_id, reviewer_compaction_hash, diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index a444bc981f..5037cc32cb 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -140,6 +140,10 @@ pub(crate) fn split_prefix_updates( } trait ErasedWorldStateSection: Send + Sync { + fn include_in_full_context(&self) -> bool { + true + } + fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool; fn has_retained_fragment_matcher(&self) -> bool; @@ -222,6 +226,10 @@ impl ErasedWorldStateSection for S { struct ExtensionWorldStateSection(WorldStateSectionContribution); impl ErasedWorldStateSection for ExtensionWorldStateSection { + fn include_in_full_context(&self) -> bool { + self.0.include_in_full_context() + } + fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool { self.0.matches_legacy_fragment(role, text) } @@ -422,6 +430,10 @@ impl From<&Map> for WorldStateSnapshot { } impl WorldStateSnapshot { + pub(crate) fn remove_section(&mut self, id: &str) { + self.sections.remove(id); + } + pub(crate) fn into_object(self) -> Map { self.sections.into_iter().collect() } @@ -499,7 +511,20 @@ impl fmt::Debug for WorldState { } } +enum RenderScope { + FullContext, + HistoryDiff, +} + impl WorldState { + pub(crate) fn full_context_excluded_section_ids( + &self, + ) -> impl Iterator + '_ { + self.sections + .iter() + .filter_map(|(id, section)| (!section.include_in_full_context()).then_some(*id)) + } + pub(crate) fn add_section(&mut self, section: S) { let id = S::ID; assert!( @@ -525,9 +550,11 @@ impl WorldState { } } - /// Renders every section as new, preserving update order and placement. + /// Renders restorable sections as new, preserving update order and placement. pub(crate) fn render_full(&self) -> (WorldStateSnapshot, Vec) { - self.render_with(/*stored*/ None, |_, _| PreviousSectionState::Absent) + self.render_with(RenderScope::FullContext, /*stored*/ None, |_, _| { + PreviousSectionState::Absent + }) } /// Falls back to retained model history when no exact persisted snapshot is available. @@ -536,7 +563,7 @@ impl WorldState { previous: Option<&WorldStateSnapshot>, items: impl IntoIterator + Clone, ) -> (WorldStateSnapshot, Vec) { - self.render_with(previous, |id, section| { + self.render_with(RenderScope::HistoryDiff, previous, |id, section| { if let Some(previous) = previous.and_then(|previous| previous.section(id)) { if section.has_retained_fragment_matcher() && !has_retained_fragment(items.clone(), section) @@ -555,12 +582,16 @@ impl WorldState { fn render_with<'a>( &self, + scope: RenderScope, stored: Option<&WorldStateSnapshot>, mut previous: impl FnMut(&str, &dyn ErasedWorldStateSection) -> PreviousSectionState<'a, Value>, ) -> (WorldStateSnapshot, Vec) { let mut snapshot = WorldStateSnapshot::default(); let mut updates = Vec::new(); for (id, section) in &self.sections { + if matches!(scope, RenderScope::FullContext) && !section.include_in_full_context() { + continue; + } let prior = previous(id, section.as_ref()); let (section_snapshot, section_updates) = section.render_diff(prior); // A skipped snapshot retains stored state even if history required a fresh render. diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 05c99d2b11..c54f109aab 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -103,6 +103,8 @@ pub(crate) async fn build_prompt_input_from_session( .await; } + sess.record_step_world_state_if_changed(step_context.as_ref()) + .await?; let prompt_input = sess .clone_history() .await diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 667618102c..8413d42879 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -482,6 +482,7 @@ pub(crate) enum CompactedHistoryInstallation<'a> { prefix: CompactionPrefix, prefix_context: Vec, initial_context: Vec, + reset_world_state_sections: Vec<&'static str>, model_info: &'a ModelInfo, compaction_response_id: String, reviewer_compaction_hash: Option, @@ -4539,6 +4540,7 @@ impl Session { prefix, prefix_context, initial_context, + reset_world_state_sections, model_info, compaction_response_id, reviewer_compaction_hash, @@ -4575,6 +4577,11 @@ impl Session { if let Some(checkpoint) = state.history.world_state_checkpoint() { world_state_baseline = WorldStateSnapshot::from(&checkpoint.state); } + // Volatile sections must refresh in the new window even if their + // previous messages and comparison state survived in the suffix. + for section_id in reset_world_state_sections { + world_state_baseline.remove_section(section_id); + } // Prefix replacement already retains all edits in the live suffix. let input_goal_ids = Default::default(); let (window_number, window_ids) = state.advance_auto_compact_window(); diff --git a/codex-rs/core/tests/suite/prompt_debug_tests.rs b/codex-rs/core/tests/suite/prompt_debug_tests.rs index c216cec8c9..37fb4d2adf 100644 --- a/codex-rs/core/tests/suite/prompt_debug_tests.rs +++ b/codex-rs/core/tests/suite/prompt_debug_tests.rs @@ -4,17 +4,48 @@ use anyhow::Result; use codex_core::build_prompt_input; use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; +use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::RenderedWorldStateFragment; +use codex_extension_api::WorldStateContributionInput; +use codex_extension_api::WorldStateSectionContribution; +use codex_extension_api::input_text; use codex_home::CodexHomeUserInstructionsProvider; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::user_input::UserInput; use core_test_support::responses::strip_metadata; use core_test_support::responses::strip_response_item_id; +use core_test_support::responses::user_message_item; use pretty_assertions::assert_eq; use tempfile::TempDir; const TEST_INSTRUCTIONS: &str = "Global test instructions"; +const STEP_REMINDER: &str = "Fresh step reminder"; + +struct StepReminder; + +impl ContextContributor for StepReminder { + fn contribute_world_state<'a>( + &'a self, + _input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async { + vec![ + WorldStateSectionContribution::new("step_reminder", |_| { + let reminder = RenderedWorldStateFragment::new( + "developer", + ("", ""), + input_text(STEP_REMINDER), + ); + (None, Some(reminder)) + }) + .without_full_context(), + ] + }) + } +} #[tokio::test] async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { @@ -33,6 +64,8 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { let user_instructions_provider = Arc::new(CodexHomeUserInstructionsProvider::new( config.codex_home.clone(), )); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.prompt_contributor(Arc::new(StepReminder)); let input = build_prompt_input( config, vec![UserInput::Text { @@ -40,31 +73,23 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { text_elements: Vec::new(), }], /*state_db*/ None, - Arc::new(ExtensionRegistryBuilder::new().build()), + Arc::new(extensions.build()), user_instructions_provider, ) .await?; - let expected_user_message = ResponseItem::Message { - status: None, - encrypted_content: None, - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "hello from debug prompt".to_string(), - }], - phase: None, - internal_chat_message_metadata_passthrough: None, + let [initial_context @ .., user_message, reminder] = input.as_slice() else { + panic!("expected user input followed by step context"); }; assert_eq!( - input - .last() - .cloned() - .map(strip_metadata) - .map(strip_response_item_id), - Some(expected_user_message) + strip_response_item_id(strip_metadata(user_message.clone())), + user_message_item("hello from debug prompt") ); - assert!(input.iter().any(|item| { + assert!( + matches!(reminder, ResponseItem::Message { role, content, .. } + if role == "developer" && content == &[input_text(STEP_REMINDER)]) + ); + let contains_text = |item: &ResponseItem, needle: &str| { let ResponseItem::Message { content, .. } = item else { return false; }; @@ -75,8 +100,18 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> { else { return false; }; - text.contains(TEST_INSTRUCTIONS) + text.contains(needle) }) - })); + }; + assert!( + initial_context + .iter() + .any(|item| contains_text(item, TEST_INSTRUCTIONS)) + ); + assert!( + !initial_context + .iter() + .any(|item| contains_text(item, STEP_REMINDER)) + ); Ok(()) } diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 530e508ea9..f974eb93ab 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -97,6 +97,7 @@ use core_test_support::stdio_server_bin; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::TestCodexBuilder; use core_test_support::test_codex::environment_config_for_selection; use core_test_support::test_codex::executor_path_uri; use core_test_support::test_codex::test_codex; @@ -228,6 +229,9 @@ mod frozen_cloud_skills; #[path = "scenarios_compaction_tests.rs"] mod compaction; +#[path = "scenarios_budget_context.rs"] +mod budget_context; + #[path = "scenarios_world_state_prefix_tests.rs"] mod world_state_prefix; @@ -3201,13 +3205,14 @@ fn compaction_output(value: &str) -> serde_json::Value { } async fn async_compaction_test( + builder: TestCodexBuilder, responses_uri: &str, features: Vec, configure: impl Fn(&mut Config) + Send + Sync + 'static, ) -> Result<(wiremock::MockServer, TestCodex)> { let config_server = start_mock_server().await; let base_url = format!("{responses_uri}/v1"); - let test = test_codex() + let test = builder .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_config(move |config| { configure_scenario_catalog(config); @@ -3437,13 +3442,23 @@ async fn astra_async_compaction_keeps_work_added_after_its_snapshot( vec![Feature::RetainClientToolOutputs] } }; - let (_config_server, test) = async_compaction_test(streaming.uri(), features, move |config| { - if over_budget { - config.async_compaction_threshold_percent = 1; - config.model_auto_compact_token_limit = Some(5_000); - } - }) - .await?; + let reminders = Arc::new(budget_context::PrefixBudgetReminders { + remaining: 400.into(), + }); + let mut builder = test_codex(); + if matches!(suffix, LiveSuffix::User) { + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.prompt_contributor(reminders.clone()); + builder = builder.with_extensions(Arc::new(extensions.build())); + } + let (_config_server, test) = + async_compaction_test(builder, streaming.uri(), features, move |config| { + if over_budget { + config.async_compaction_threshold_percent = 1; + config.model_auto_compact_token_limit = Some(5_000); + } + }) + .await?; test.codex.ensure_rollout_materialized().await; let compaction_event = |event: &EventMsg| match event { EventMsg::ItemStarted(event) if matches!(event.item, TurnItem::ContextCompaction(_)) => { @@ -3459,6 +3474,9 @@ async fn astra_async_compaction_keeps_work_added_after_its_snapshot( let timeout = Duration::from_secs(/*secs*/ 20); tokio::time::timeout(timeout, streaming.wait_for_request_count(/*count*/ 2)).await?; let next_input = if matches!(suffix, LiveSuffix::User) { + reminders + .remaining + .store(350, std::sync::atomic::Ordering::SeqCst); TurnInputRequest::user_input(vec![text("Check who can finish the notes.")]) } else { let item = serde_json::from_value(json!({ @@ -3516,6 +3534,9 @@ async fn astra_async_compaction_keeps_work_added_after_its_snapshot( test.codex.wait_for_background_compaction_for_test(), ) .await?; + reminders + .remaining + .store(325, std::sync::atomic::Ordering::SeqCst); release_foreground .send(()) .expect("release foreground work"); @@ -3567,6 +3588,8 @@ async fn astra_async_compaction_keeps_work_added_after_its_snapshot( assert!(history.windows(2).any(|items| matches!(items, [RolloutItem::Compacted(_), RolloutItem::WorldState(baseline)] if baseline.full && baseline.state["collaboration_mode"]["mode"] == "plan" + && !baseline.state.contains_key("cutoff_budget") + && !baseline.state.contains_key("install_budget") ))); let input = requests[3]["input"].as_array().expect("request input"); // The full snapshot normalizes mode text; check the cutoff and later update explicitly. @@ -3590,7 +3613,7 @@ async fn astra_async_compaction_keeps_work_added_after_its_snapshot( insta::assert_snapshot!( "astra_async_compaction_with_live_suffix", context_snapshot::format_context_snapshot( - "Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode while it runs. Adoption restores Default mode at the cutoff, then preserves the Plan update, newer user message, and assistant/tool exchange in their original order.", + "Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode and the host changes which budget reminder opts out of full context. Adoption keeps the live suffix and refreshes both reminders after the tool output.", &entries, &ContextSnapshotOptions::default().rewrite_known_segments(), ) @@ -3715,6 +3738,7 @@ async fn async_compaction_reuses_pending_work( ))); let (streaming, mut completions) = start_streaming_sse_server(responses).await; let (_config_server, test) = async_compaction_test( + test_codex(), streaming.uri(), if over_budget { vec![Feature::RetainClientToolOutputs] diff --git a/codex-rs/core/tests/suite/scenarios_budget_context.rs b/codex-rs/core/tests/suite/scenarios_budget_context.rs new file mode 100644 index 0000000000..433eb006a6 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_budget_context.rs @@ -0,0 +1,224 @@ +//! Host budget reminders use generic World State and refresh after compaction spending. + +use super::*; +use codex_extension_api::ContextContributor; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::PreviousWorldStateSection; +use codex_extension_api::RenderedWorldStateFragment; +use codex_extension_api::WorldStateContributionInput; +use codex_extension_api::WorldStateSectionContribution; +use codex_extension_api::input_text; +use core_test_support::responses::received_responses_requests; +use core_test_support::responses::sse_response; +use core_test_support::responses::strip_metadata_from_json; +use core_test_support::responses::strip_response_item_ids_from_json; +use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; +use wiremock::Mock; +use wiremock::matchers::method; +use wiremock::matchers::path_regex; + +struct HostBudgetContext { + remaining: Arc, +} + +impl ContextContributor for HostBudgetContext { + fn contribute_world_state<'a>( + &'a self, + _input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + let remaining = self.remaining.load(Ordering::SeqCst); + let threshold = [750, 400, 100] + .into_iter() + .filter(|threshold| remaining <= *threshold) + .count(); + Box::pin(async move { + vec![ + WorldStateSectionContribution::new("host_budget_reminder", move |previous| { + if let PreviousWorldStateSection::Known(previous) = previous + && previous.as_u64().is_some_and(|index| index >= threshold as u64) + { + return (None, None); + } + ( + Some(json!(threshold)), + Some(RenderedWorldStateFragment::new( + "developer", + ("\n", "\n"), + input_text(format!( + "You have {remaining} weighted tokens left in the shared session token budget." + )), + )), + ) + }) + .with_retained_fragment_matcher(|role, text| { + role == "developer" && text.contains("") + }) + .without_full_context(), + ] + }) + } +} + +// At 400 the cutoff section opts out. At 350 and 325 the install section opts out. +// Both later values share a comparison bucket, so rendering 325 proves the markers were cleared. +pub(super) struct PrefixBudgetReminders { + pub(super) remaining: AtomicI64, +} + +impl ContextContributor for PrefixBudgetReminders { + fn contribute_world_state<'a>( + &'a self, + _input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + let remaining = self.remaining.load(Ordering::SeqCst); + let section = move |id: &'static str| { + WorldStateSectionContribution::new(id, move |previous| { + let threshold = json!(remaining / 100); + if matches!(previous, PreviousWorldStateSection::Known(value) if value == &threshold) + { + return (None, None); + } + ( + Some(threshold), + Some(RenderedWorldStateFragment::new( + "developer", + ("", ""), + input_text(format!("{id}: {remaining} tokens remain.")), + )), + ) + }) + }; + Box::pin(async move { + let cutoff = section("cutoff_budget"); + let install = section("install_budget"); + if remaining == 400 { + vec![cutoff.without_full_context(), install] + } else { + vec![cutoff, install.without_full_context()] + } + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + // This fixture supplies host-owned policy and controlled shared spending. It does + // not implement usage reporting or configure Local's native rollout budget. + let remaining = Arc::new(AtomicI64::new(1_000)); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.prompt_contributor(Arc::new(HostBudgetContext { + remaining: Arc::clone(&remaining), + })); + let test = test_codex() + .with_model("gpt-6-astra") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + configure_scenario_catalog(config); + config.workspace_roots = vec![config.cwd.clone()]; + config.rollout_budget = None; + }) + .build_with_auto_env(&server) + .await?; + + mount_sse_sequence( + &server, + vec![ + sse(vec![ev_completed("initial")]), + sse(vec![ev_completed("same-threshold")]), + sse(vec![ev_completed("crossed-threshold")]), + ], + ) + .await; + test.submit_turn("Start planning the release.").await?; + remaining.store(900, Ordering::SeqCst); + test.submit_turn("Continue planning.").await?; + remaining.store(400, Ordering::SeqCst); + test.submit_turn("Check the release plan.").await?; + + let compaction_budget = Arc::clone(&remaining); + Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .respond_with(move |_: &wiremock::Request| { + // World State was already captured at 400 before this request. Simulate + // compaction spending before Core installs that captured snapshot. + compaction_budget.store(350, Ordering::SeqCst); + sse_response(sse(vec![ + json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction", + "encrypted_content": "BUDGET_CONTEXT_CHECKPOINT", + }, + }), + ev_completed("compact-response"), + ])) + }) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + assert!(!matches!(event, EventMsg::Error(_)), "{event:?}"); + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + mount_sse_sequence( + &server, + vec![ + sse(vec![ev_completed("fresh-window")]), + sse(vec![ev_completed("same-fresh-threshold")]), + ], + ) + .await; + test.submit_turn("Continue from the compacted plan.") + .await?; + remaining.store(300, Ordering::SeqCst); + test.submit_turn("Finish the release plan.").await?; + + let requests = received_responses_requests(&server).await; + assert_eq!(requests.len(), 6); + // Compare the actual wire input without normalization: both unchanged and + // threshold-crossing reminders must only append within a context window. + for (previous, next) in [(0, 1), (1, 2), (4, 5)] { + let previous_input = requests[previous].input(); + let next_input = requests[next].input(); + assert_eq!( + next_input.get(..previous_input.len()), + Some(previous_input.as_slice()), + "request {next} must preserve request {previous}'s entire input prefix", + ); + for field in ["prompt_cache_key", "instructions", "tools"] { + assert_eq!( + requests[next].body_json()[field], + requests[previous].body_json()[field], + "request {next} must preserve request {previous}'s {field}", + ); + } + } + let resumed_input = requests[4].input(); + // Compaction rebuilds item IDs and metadata while preserving the tool + // declarations and base instructions ahead of the conversation. + let initial_prefix = serde_json::Value::Array(requests[0].input()[..2].to_vec()); + let resumed_prefix = serde_json::Value::Array(resumed_input[..2].to_vec()); + assert_eq!( + strip_metadata_from_json(strip_response_item_ids_from_json(initial_prefix)), + strip_metadata_from_json(strip_response_item_ids_from_json(resumed_prefix)), + ); + insta::assert_snapshot!( + "host_budget_context_after_compaction", + context_snapshot::format_request_history_snapshot( + "A host contributes budget reminders at thresholds. Remote compaction spends within the same threshold, so its captured reminder is omitted and the next sampling step inserts the fresh remainder once.", + &requests, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ) + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_async_compaction_with_live_suffix.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_async_compaction_with_live_suffix.snap index b13ef69cfa..6d5fb853a3 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_async_compaction_with_live_suffix.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__astra_async_compaction_with_live_suffix.snap @@ -1,8 +1,8 @@ --- source: core/tests/suite/scenarios.rs -expression: "context_snapshot::format_context_snapshot(\"Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode while it runs. Adoption restores Default mode at the cutoff, then preserves the Plan update, newer user message, and assistant/tool exchange in their original order.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +expression: "context_snapshot::format_context_snapshot(\"Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode and the host changes which budget reminder opts out of full context. Adoption keeps the live suffix and refreshes both reminders after the tool output.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" --- -Scenario: Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode while it runs. Adoption restores Default mode at the cutoff, then preserves the Plan update, newer user message, and assistant/tool exchange in their original order. +Scenario: Astra compacts a completed launch-notes turn in the background. The user switches to Plan mode and the host changes which budget reminder opts out of full context. Adoption keeps the live suffix and refreshes both reminders after the tool output. ## Window 1 -- request 1 (request) -- @@ -23,9 +23,10 @@ Scenario: Astra compacts a completed launch-notes turn in the background. The us - function/wait_agent 01:message/developer: -02:message/developer[2]: +02:message/developer[3]: [01] [02] + [03] install_budget: 400 tokens remain. 03:message/developer: 04:message/developer: @@ -34,13 +35,15 @@ Scenario: Astra compacts a completed launch-notes turn in the background. The us 06:message/user: Draft the launch notes. -07:configuration_update:reasoning={"effort":"low"} +07:message/developer: + cutoff_budget: 400 tokens remain. +08:configuration_update:reasoning={"effort":"low"} -- request 2 (request) -- -08:message/assistant: +09:message/assistant: I drafted the launch notes. -09:compaction_trigger +10:compaction_trigger -## Window 2 (after request 2: input diverged at item 09) +## Window 2 (after request 2: input diverged at item 10) -- request 3 (request) -- 00:additional_tools/developer (3; hash=bf2e30605ff5dc17): - namespace/functions; hash=fada4a476a3754dd @@ -59,9 +62,10 @@ Scenario: Astra compacts a completed launch-notes turn in the background. The us - function/wait_agent 01:message/developer: -02:message/developer[2]: +02:message/developer[3]: [01] [02] + [03] install_budget: 400 tokens remain. 03:message/developer: 04:message/developer: @@ -70,12 +74,16 @@ Scenario: Astra compacts a completed launch-notes turn in the background. The us 06:message/user: Draft the launch notes. -07:configuration_update:reasoning={"effort":"low"} -08:message/assistant: +07:message/developer: + cutoff_budget: 400 tokens remain. +08:configuration_update:reasoning={"effort":"low"} +09:message/assistant: I drafted the launch notes. -09:message/developer: - -10:message/user: +10:message/developer[3]: + [01] + [02] cutoff_budget: 350 tokens remain. + [03] install_budget: 350 tokens remain. +11:message/user: Check who can finish the notes. ## Window 3 (after request 3: input diverged at item 02) @@ -100,20 +108,26 @@ Scenario: Astra compacts a completed launch-notes turn in the background. The us 02:message/user: Draft the launch notes. 03:compaction:encrypted=true; chars=23; hash=159cfcb231cc70f0 -04:message/developer[2]: +04:message/developer[3]: [01] [02] + [03] install_budget: 400 tokens remain. 05:message/developer: 06:message/developer: 07:message/user: -08:message/developer: - +08:message/developer[3]: + [01] + [02] cutoff_budget: 350 tokens remain. + [03] install_budget: 350 tokens remain. 09:message/user: Check who can finish the notes. 10:message/assistant: I will check who is available. 11:function_call/collaboration.list_agents:{} 12:function_call_output:{"agents":[{"agent_name":"/root","agent_status":"running"}]} +13:message/developer[2]: + [01] cutoff_budget: 325 tokens remain. + [02] install_budget: 325 tokens remain. diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap new file mode 100644 index 0000000000..90fc2118ea --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap @@ -0,0 +1,98 @@ +--- +source: core/tests/suite/scenarios_budget_context.rs +expression: "context_snapshot::format_request_history_snapshot(\"A host contributes budget reminders at thresholds. Remote compaction spends within the same threshold, so its captured reminder is omitted and the next sampling step inserts the fresh remainder once.\",\n&requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: A host contributes budget reminders at thresholds. Remote compaction spends within the same threshold, so its captured reminder is omitted and the next sampling step inserts the fresh remainder once. + +## Window 1 +-- request 1 (turn) -- +00:additional_tools/developer (3; hash=bf2e30605ff5dc17): + - namespace/functions; hash=fada4a476a3754dd + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63 + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent +01:message/developer: + +02:message/developer[2]: + [01] + [02] +03:message/developer: + +04:message/developer: + +05:message/user: + +06:message/user: + Start planning the release. +07:message/developer: + + You have 1000 weighted tokens left in the shared session token budget. + +08:configuration_update:reasoning={"effort":"low"} +-- request 2 (turn) -- +09:message/user: + Continue planning. +-- request 3 (turn) -- +10:message/developer: + + You have 400 weighted tokens left in the shared session token budget. + +11:message/user: + Check the release plan. +-- request 4 (compaction) -- +12:compaction_trigger + +## Window 2 (after request 4: input diverged at item 02) +-- request 5 (turn) -- +00:additional_tools/developer (3; hash=bf2e30605ff5dc17): + - namespace/functions; hash=fada4a476a3754dd + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63 + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent +01:message/developer: + +02:message/user: + Start planning the release. +03:message/user: + Continue planning. +04:message/developer[2]: + [01] + [02] +05:message/developer: + +06:message/developer: + +07:message/user: + +08:message/user: + Check the release plan. +09:compaction:encrypted=true; chars=25; hash=455278c90f6b4d93 +10:message/developer: + + You have 350 weighted tokens left in the shared session token budget. + +11:message/user: + Continue from the compacted plan. +-- request 6 (turn) -- +12:message/user: + Finish the release plan. diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index 7f30e074ad..d3bfb8b13f 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -98,6 +98,7 @@ type LegacyFragmentMatcher = dyn Fn(&str, &str) -> bool + Send + Sync; #[derive(Clone)] pub struct WorldStateSectionContribution { id: &'static str, + include_in_full_context: bool, window_prefix: bool, render_diff: Arc, matches_legacy_fragment: Arc, @@ -119,6 +120,7 @@ impl WorldStateSectionContribution { ) -> Self { Self { id, + include_in_full_context: true, window_prefix: false, render_diff: Arc::new(render_diff), matches_legacy_fragment: Arc::new(|_, _| false), @@ -126,6 +128,20 @@ impl WorldStateSectionContribution { } } + /// Recompute this section during ordinary step context updates instead of replaying + /// its captured value when initial or compacted context is rebuilt. Full rendering + /// omits both its fragment and comparison snapshot so a fresh contribution can emit. + /// After prefix compaction, the next diff receives [`PreviousWorldStateSection::Unknown`] + /// if the section's legacy matcher recognizes text in the retained suffix. + pub fn without_full_context(mut self) -> Self { + self.include_in_full_context = false; + self + } + + pub fn include_in_full_context(&self) -> bool { + self.include_in_full_context + } + pub fn with_legacy_matcher( mut self, matcher: impl Fn(&str, &str) -> bool + Send + Sync + 'static,