openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-22 17:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

core: use current step state after compaction

Browse Source
This commit is contained in:
Sayan Sisodiya
2026-06-22 19:18:07 -07:00
parent fbd575ab4a
commit b680bfccee
8 changed files with 215 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/compact.rs
+25 -4
View File
@@ -4,6 +4,7 @@ use std::time::Instant;
use crate::Prompt;
use crate::client::ModelClientSession;
use crate::client_common::ResponseEvent;
use crate::context::world_state::WorldState;
use crate::hook_runtime::PostCompactHookOutcome;
use crate::hook_runtime::PreCompactHookOutcome;
use crate::hook_runtime::run_post_compact_hooks;
@@ -66,6 +67,23 @@ pub(crate) enum InitialContextInjection {
DoNotInject,
}
pub(crate) async fn build_compaction_initial_context(
sess: &Session,
turn_context: &TurnContext,
initial_context_injection: InitialContextInjection,
) -> (Vec<ResponseItem>, Option<WorldState>) {
// Return the rendered state with its items so history and its baseline stay identical.
match initial_context_injection {
InitialContextInjection::BeforeLastUserMessage => {
let (items, world_state) = sess
.build_initial_context_for_current_step(turn_context)
.await;
(items, Some(world_state))
}
InitialContextInjection::DoNotInject => (Vec::new(), None),
}
}
pub(crate) fn should_use_remote_compact_task(provider: &ModelProviderInfo) -> bool {
provider.supports_remote_compaction()
}
@@ -309,11 +327,13 @@ async fn run_compact_task_inner_impl(
}
let (window_number, window_ids) = sess.advance_auto_compact_window().await;
if matches!(
let (initial_context, world_state_baseline) = build_compaction_initial_context(
sess.as_ref(),
turn_context.as_ref(),
initial_context_injection,
InitialContextInjection::BeforeLastUserMessage
) {
let initial_context = sess.build_initial_context(turn_context.as_ref()).await;
)
.await;
if !initial_context.is_empty() {
new_history =
insert_initial_context_before_last_real_user_or_summary(new_history, initial_context);
}
@@ -333,6 +353,7 @@ async fn run_compact_task_inner_impl(
turn_context.as_ref(),
new_history,
reference_context_item,
world_state_baseline,
compacted_item,
)
.await;
codex-rs/core/src/compact_remote.rs
+12 -12
View File
@@ -6,8 +6,10 @@ use crate::client::CompactConversationRequestSettings;
use crate::compact::CompactionAnalyticsAttempt;
use crate::compact::CompactionAnalyticsDetails;
use crate::compact::InitialContextInjection;
use crate::compact::build_compaction_initial_context;
use crate::compact::compaction_status_from_result;
use crate::compact::insert_initial_context_before_last_real_user_or_summary;
use crate::context::world_state::WorldState;
use crate::context_manager::ContextManager;
use crate::hook_runtime::PostCompactHookOutcome;
use crate::hook_runtime::PreCompactHookOutcome;
@@ -237,7 +239,7 @@ async fn run_remote_compact_task_inner_impl(
window_id,
CodexResponsesRequestKind::Compaction(compaction_metadata),
);
let mut new_history = sess
let new_history = sess
.services
.model_client
.compact_conversation_history(
@@ -259,7 +261,7 @@ async fn run_remote_compact_task_inner_impl(
)
.await?;
let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await;
new_history = process_compacted_history(
let (new_history, world_state_baseline) = process_compacted_history(
sess.as_ref(),
turn_context.as_ref(),
new_history,
@@ -290,6 +292,7 @@ async fn run_remote_compact_task_inner_impl(
turn_context.as_ref(),
new_history,
reference_context_item,
world_state_baseline,
compacted_item,
)
.await;
@@ -305,21 +308,18 @@ pub(crate) async fn process_compacted_history(
turn_context: &TurnContext,
mut compacted_history: Vec<ResponseItem>,
initial_context_injection: InitialContextInjection,
) -> Vec<ResponseItem> {
) -> (Vec<ResponseItem>, Option<WorldState>) {
// Mid-turn compaction is the only path that must inject initial context above the last user
// message in the replacement history. Pre-turn compaction instead injects context after the
// compaction item, but mid-turn compaction keeps the compaction item last for model training.
let initial_context = if matches!(
initial_context_injection,
InitialContextInjection::BeforeLastUserMessage
) {
sess.build_initial_context(turn_context).await
} else {
Vec::new()
};
let (initial_context, world_state_baseline) =
build_compaction_initial_context(sess, turn_context, initial_context_injection).await;
compacted_history.retain(should_keep_compacted_history_item);
insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context)
(
insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context),
world_state_baseline,
)
}
/// Returns whether an item from remote compaction output should be preserved.
codex-rs/core/src/compact_remote_v2.rs
+2 -1
View File
@@ -293,7 +293,7 @@ async fn run_remote_compact_task_inner_impl(
build_v2_compacted_history(&prompt_input, compaction_output);
analytics_details.retained_image_count = Some(retained_images);
let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await;
let new_history = process_compacted_history(
let (new_history, world_state_baseline) = process_compacted_history(
sess.as_ref(),
turn_context.as_ref(),
compacted_history,
@@ -321,6 +321,7 @@ async fn run_remote_compact_task_inner_impl(
turn_context.as_ref(),
new_history,
reference_context_item,
world_state_baseline,
compacted_item,
)
.await;
codex-rs/core/src/compact_tests.rs
+1 -1
View File
@@ -14,7 +14,7 @@ async fn process_compacted_history_with_test_session(
.set_previous_turn_settings(previous_turn_settings.cloned())
.await;
let initial_context = session.build_initial_context(&turn_context).await;
let refreshed = crate::compact_remote::process_compacted_history(
let (refreshed, _) = crate::compact_remote::process_compacted_history(
&session,
&turn_context,
compacted_history,
codex-rs/core/src/session/mod.rs
+45 -8
View File
@@ -32,6 +32,7 @@ use crate::context::NetworkRuleSaved;
use crate::context::PermissionsInstructions;
use crate::context::PersonalitySpecInstructions;
use crate::context::RecommendedPluginsInstructions;
use crate::context::world_state::EnvironmentsState;
use crate::context::world_state::WorldState;
use crate::current_time::TimeProvider;
use crate::default_skill_metadata_budget;
@@ -238,6 +239,7 @@ use self::session::AppServerClientMetadata;
use self::session::Session;
use self::session::SessionConfiguration;
pub(crate) use self::session::SessionSettingsUpdate;
use self::step_context::StepContext;
#[cfg(test)]
use self::turn::AssistantMessageStreamParsers;
#[cfg(test)]
@@ -2785,6 +2787,19 @@ impl Session {
self.record_conversation_items(turn_context, &items).await;
}
pub(crate) async fn capture_step_context(&self, turn_context: &TurnContext) -> StepContext {
let environments = if turn_context
.config
.features
.enabled(Feature::DeferredExecutor)
{
self.services.turn_environments.snapshot().await
} else {
turn_context.environments.clone()
};
StepContext { environments }
}
pub(crate) async fn record_inter_agent_communication(
&self,
turn_context: &TurnContext,
@@ -2884,6 +2899,7 @@ impl Session {
turn_context: &TurnContext,
items: Vec<ResponseItem>,
reference_context_item: Option<TurnContextItem>,
world_state_baseline: Option<WorldState>,
compacted_item: CompactedItem,
) {
let items = if turn_context.config.features.enabled(Feature::ItemIds) {
@@ -2895,11 +2911,6 @@ impl Session {
replacement_history: Some(items.clone()),
..compacted_item
};
let world_state_baseline = if reference_context_item.is_some() {
Some(self.build_world_state(turn_context).await)
} else {
None
};
{
let mut state = self.state.lock().await;
state.replace_history(items, reference_context_item.clone());
@@ -3031,6 +3042,7 @@ impl Session {
items
}
#[cfg(test)]
pub(crate) async fn build_initial_context(
&self,
turn_context: &TurnContext,
@@ -3040,6 +3052,32 @@ impl Session {
.await
}
pub(crate) async fn build_initial_context_for_current_step(
&self,
turn_context: &TurnContext,
) -> (Vec<ResponseItem>, WorldState) {
let step_context = self.capture_step_context(turn_context).await;
let mut world_state = WorldState::default();
if turn_context.config.include_environment_context {
let environment_subagents = self
.services
.agent_control
.format_environment_context_subagents(self.thread_id)
.await;
world_state.add_section(
EnvironmentsState::from_turn_context_with_environments(
turn_context,
&step_context.environments,
)
.with_subagents(environment_subagents),
);
}
let items = self
.build_initial_context_with_world_state(turn_context, &world_state)
.await;
(items, world_state)
}
async fn build_world_state(&self, turn_context: &TurnContext) -> WorldState {
let environment_subagents = if turn_context.config.include_environment_context {
self.services
@@ -3407,9 +3445,8 @@ impl Session {
state.start_new_context_window_if_requested()
};
let (window_number, window_ids) = window?;
let world_state = self.build_world_state(turn_context).await;
let context_items = self
.build_initial_context_with_world_state(turn_context, &world_state)
let (context_items, world_state) = self
.build_initial_context_for_current_step(turn_context)
.await;
let turn_context_item = turn_context.to_turn_context_item();
let replacement_history = context_items;
codex-rs/core/src/session/tests/guardian_tests.rs
+1 -1
View File
@@ -524,7 +524,7 @@ async fn process_compacted_history_preserves_separate_guardian_developer_message
turn_context.session_source = guardian_source;
turn_context.developer_instructions = Some(guardian_policy.clone());
let refreshed = crate::compact_remote::process_compacted_history(
let (refreshed, _) = crate::compact_remote::process_compacted_history(
&session,
&turn_context,
vec![
codex-rs/core/src/session/turn.rs
+1 -4
View File
@@ -42,7 +42,6 @@ use crate::responses_retry::handle_retryable_response_stream_error;
use crate::session::PreviousTurnSettings;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::stream_events_utils::HandleOutputCtx;
use crate::stream_events_utils::TurnItemContributorPolicy;
@@ -247,9 +246,7 @@ pub(crate) async fn run_turn(
.features
.enabled(Feature::DeferredExecutor)
{
let step_context = StepContext {
environments: sess.services.turn_environments.snapshot().await,
};
let step_context = sess.capture_step_context(turn_context.as_ref()).await;
sess.record_step_environment_context_if_changed(
turn_context.as_ref(),
&step_context,
codex-rs/core/tests/suite/remote_env.rs
+128
View File
@@ -1,6 +1,7 @@
use anyhow::Context;
use anyhow::Result;
use codex_config::types::ApprovalsReviewer;
use codex_core::compact::SUMMARIZATION_PROMPT;
use codex_core::config::Constrained;
use codex_exec_server::CopyOptions;
use codex_exec_server::CreateDirectoryOptions;
@@ -38,6 +39,7 @@ use core_test_support::get_remote_test_env;
use core_test_support::responses::ev_apply_patch_custom_tool_call;
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_sse_once;
@@ -683,6 +685,132 @@ async fn deferred_executor_updates_model_context_after_startup() -> Result<()> {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn deferred_executor_compaction_reinjects_current_environment_once() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let server = start_mock_server().await;
let response_mock = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created("resp-1"),
ev_function_call(
"wait-for-startup",
"request_user_input",
&json!({
"questions": [{
"id": "continue",
"header": "Continue",
"question": "Continue after startup?",
"options": [{
"label": "Yes (Recommended)",
"description": "Continue the test."
}, {
"label": "No",
"description": "Stop the test."
}]
}]
})
.to_string(),
),
ev_completed_with_tokens("resp-1", /*total_tokens*/ 96),
]),
sse(vec![
ev_assistant_message("msg-compact", "AUTO_COMPACT_SUMMARY"),
ev_completed_with_tokens("resp-compact", /*total_tokens*/ 10),
]),
sse(vec![
ev_response_created("resp-2"),
ev_assistant_message("msg-2", "done"),
ev_completed("resp-2"),
]),
],
)
.await;
let mut builder = test_codex()
.with_exec_server_url(format!("ws://{}", listener.local_addr()?))
.with_config(|config| {
assert!(config.features.enable(Feature::DeferredExecutor).is_ok());
assert!(
config
.features
.enable(Feature::DefaultModeRequestUserInput)
.is_ok()
);
config.model_provider.name = "OpenAI (test)".to_string();
config.compact_prompt = Some(SUMMARIZATION_PROMPT.to_string());
config.model_context_window = Some(100);
config.model_auto_compact_token_limit = Some(90);
});
let test = timeout(Duration::from_secs(5), builder.build(&server))
.await
.context("thread startup should not wait for the remote environment")??;
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "wait for the environment".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
let request = wait_for_event_match(&test.codex, |event| match event {
EventMsg::RequestUserInput(request) => Some(request.clone()),
_ => None,
})
.await;
serve_environment_info(listener).await;
test.codex
.submit(Op::UserInputAnswer {
id: request.turn_id,
response: RequestUserInputResponse {
answers: HashMap::from([(
"continue".to_string(),
RequestUserInputAnswer {
answers: vec!["Yes (Recommended)".to_string()],
},
)]),
},
})
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let requests = response_mock.requests();
assert_eq!(requests.len(), 3);
let initial_context = requests[0].message_input_texts("user");
assert!(
initial_context
.iter()
.any(|text| text.contains("<status>starting</status>"))
);
let post_compaction_context = requests[2].message_input_texts("user");
assert_eq!(
post_compaction_context
.iter()
.filter(|text| text.contains("<status>starting</status>"))
.count(),
0
);
assert_eq!(
post_compaction_context
.iter()
.filter(|text| text.contains("<shell>zsh</shell>"))
.count(),
1
);
Ok(())
}
fn absolute_path(path: PathBuf) -> AbsolutePathBuf {
AbsolutePathBuf::try_from(path).expect("path should be absolute")
}