openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-03 19:53:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

harden codex resume fast paths

Browse Source
This commit is contained in:
anais
2026-06-21 11:16:33 -07:00
parent 661646a6b2
commit a95cf32abb
2 changed files with 259 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors/thread_processor.rs
+112 -34
View File
@@ -4263,6 +4263,8 @@ fn read_compacted_resume_prefix_items(
let file = std::fs::File::open(path)?;
let reader = std::io::BufReader::new(file.take(byte_len));
let mut prefix_items = Vec::with_capacity(3);
let mut latest_metadata_candidates = Vec::new();
let mut latest_metadata_before_checkpoint = Vec::new();
let mut saw_session_meta = false;
let mut saw_user_event = false;
let mut saw_user_response_item = false;
@@ -4284,33 +4286,54 @@ fn read_compacted_resume_prefix_items(
Ok(rollout_line) => rollout_line,
Err(_) => return Ok(None),
};
let is_reconstruction_checkpoint =
rollout_item_is_reconstruction_checkpoint(&rollout_line.item);
match &rollout_line.item {
RolloutItem::SessionMeta(meta_line)
if meta_line.meta.id == expected_thread_id && !saw_session_meta =>
{
saw_session_meta = true;
prefix_items.push(rollout_line.item);
prefix_items.push(rollout_line.item.clone());
}
RolloutItem::SessionMeta(_) => return Ok(None),
RolloutItem::EventMsg(EventMsg::UserMessage(_)) if !saw_user_event => {
saw_user_event = true;
prefix_items.push(rollout_line.item);
prefix_items.push(rollout_line.item.clone());
}
RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::Message {
role,
..
}) if role == "user" && !saw_user_response_item => {
saw_user_response_item = true;
prefix_items.push(rollout_line.item);
prefix_items.push(rollout_line.item.clone());
}
item if rollout_item_needs_full_resume_load(item) => {
return Ok(None);
}
RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_)) => return Ok(None),
RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(_)) => {
latest_metadata_candidates.retain(|item| {
!matches!(item, RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(_)))
});
latest_metadata_candidates.push(rollout_line.item.clone());
}
RolloutItem::EventMsg(EventMsg::TokenCount(_)) => {
latest_metadata_candidates
.retain(|item| !matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))));
latest_metadata_candidates.push(rollout_line.item.clone());
}
item if rollout_item_needs_full_resume_load(item) => return Ok(None),
_ => {}
}
if saw_session_meta && saw_user_event && saw_user_response_item {
return Ok(Some(prefix_items));
if is_reconstruction_checkpoint {
latest_metadata_before_checkpoint = latest_metadata_candidates.clone();
}
}
Ok(saw_session_meta.then_some(prefix_items))
if saw_session_meta {
prefix_items.extend(latest_metadata_before_checkpoint);
Ok(Some(prefix_items))
} else {
Ok(None)
}
}
fn read_compacted_resume_suffix_items(
@@ -4320,6 +4343,7 @@ fn read_compacted_resume_suffix_items(
let mut suffix_items = Vec::new();
let mut found_checkpoint = false;
let mut found_checkpoint_turn_context = false;
let mut found_checkpoint_user_boundary = false;
let mut should_fallback = false;
read_rollout_lines_from_end(path, byte_len, |line| {
@@ -4335,12 +4359,29 @@ fn read_compacted_resume_suffix_items(
return Ok(false);
}
};
if found_checkpoint
&& matches!(
&rollout_line.item,
RolloutItem::EventMsg(
EventMsg::ThreadRolledBack(_)
| EventMsg::ThreadGoalUpdated(_)
| EventMsg::TokenCount(_)
)
)
{
should_fallback = true;
return Ok(false);
}
if rollout_item_needs_full_resume_load(&rollout_line.item) {
should_fallback = true;
return Ok(false);
}
let is_checkpoint = rollout_item_is_reconstruction_checkpoint(&rollout_line.item);
let is_turn_context = matches!(&rollout_line.item, RolloutItem::TurnContext(_));
let is_user_boundary = matches!(
&rollout_line.item,
RolloutItem::EventMsg(EventMsg::UserMessage(_))
);
suffix_items.push(rollout_line.item);
if is_checkpoint {
found_checkpoint = true;
@@ -4348,12 +4389,20 @@ fn read_compacted_resume_suffix_items(
}
if found_checkpoint && is_turn_context {
found_checkpoint_turn_context = true;
return Ok(true);
}
if found_checkpoint && is_user_boundary {
found_checkpoint_user_boundary = true;
return Ok(false);
}
Ok(true)
})?;
if should_fallback || !found_checkpoint || !found_checkpoint_turn_context {
if should_fallback
|| !found_checkpoint
|| !found_checkpoint_turn_context
|| !found_checkpoint_user_boundary
{
return Ok(None);
}
@@ -4433,60 +4482,89 @@ fn rollout_item_needs_full_resume_load(item: &RolloutItem) -> bool {
}
}
const RECENT_NOT_LOADED_SUFFIX_SCAN_LIMIT_BYTES: usize = 2 * 1024 * 1024;
fn read_recent_not_loaded_rollout_items(
path: &Path,
page_size: usize,
) -> std::io::Result<Option<RecentRolloutItems>> {
let contents = std::fs::read_to_string(path)?;
if contents
.lines()
.find(|line| !line.trim().is_empty())
.is_some_and(rollout_line_has_fork_parent_rollout_ref)
{
let metadata = std::fs::metadata(path)?;
if !metadata.is_file() || metadata.len() == 0 {
return Ok(None);
}
if rollout_head_has_fork_parent_rollout_ref(path)? {
return Ok(None);
}
let desired_turns = page_size.saturating_add(1);
let mut record_index_after = contents
.lines()
.filter(|line| !line.trim().is_empty())
.count();
let mut suffix_start_index = record_index_after;
let mut suffix_items = Vec::new();
let mut turn_openings = 0usize;
let mut explicit_turn_starts = 0usize;
let mut scanned_bytes = 0usize;
let mut should_fallback = false;
for line in contents.lines().rev() {
read_rollout_lines_from_end(path, metadata.len(), |line| {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
return Ok(true);
}
scanned_bytes = scanned_bytes.saturating_add(line.len()).saturating_add(1);
if scanned_bytes > RECENT_NOT_LOADED_SUFFIX_SCAN_LIMIT_BYTES
&& explicit_turn_starts < desired_turns
{
should_fallback = true;
return Ok(false);
}
record_index_after = record_index_after.saturating_sub(1);
let rollout_line =
match serde_json::from_str::<codex_protocol::protocol::RolloutLine>(trimmed) {
Ok(rollout_line) => rollout_line,
Err(_) => return Ok(None),
Err(_) => {
should_fallback = true;
return Ok(false);
}
};
if rollout_item_disables_recent_not_loaded_fast_path(&rollout_line.item) {
return Ok(None);
if !is_persisted_rollout_item(&rollout_line.item)
|| rollout_item_disables_recent_not_loaded_fast_path(&rollout_line.item)
{
should_fallback = true;
return Ok(false);
}
if rollout_item_opens_implicit_turn(&rollout_line.item) {
turn_openings = turn_openings.saturating_add(1);
if rollout_item_starts_explicit_turn(&rollout_line.item) {
explicit_turn_starts = explicit_turn_starts.saturating_add(1);
}
suffix_start_index = record_index_after;
suffix_items.push(rollout_line.item);
if turn_openings >= desired_turns {
break;
if explicit_turn_starts >= desired_turns {
return Ok(false);
}
Ok(true)
})?;
if should_fallback || explicit_turn_starts == 0 {
return Ok(None);
}
suffix_items.reverse();
Ok(Some(RecentRolloutItems {
initial_rollout_index: suffix_start_index,
initial_rollout_index: 0,
items: suffix_items,
}))
}
fn rollout_item_opens_implicit_turn(item: &RolloutItem) -> bool {
matches!(item, RolloutItem::EventMsg(EventMsg::UserMessage(_)))
fn rollout_head_has_fork_parent_rollout_ref(path: &Path) -> std::io::Result<bool> {
use std::io::BufRead;
let file = std::fs::File::open(path)?;
let reader = std::io::BufReader::new(file);
for line in reader.lines() {
let line = line?;
let trimmed = line.trim();
if !trimmed.is_empty() {
return Ok(rollout_line_has_fork_parent_rollout_ref(trimmed));
}
}
Ok(false)
}
fn rollout_item_starts_explicit_turn(item: &RolloutItem) -> bool {
matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(_)))
}
fn rollout_item_disables_recent_not_loaded_fast_path(item: &RolloutItem) -> bool {
codex-rs/app-server/src/request_processors/thread_processor_tests.rs
+147 -4
View File
@@ -416,6 +416,34 @@ mod thread_processor_behavior_tests {
})
}
fn compacted_resume_test_token_count() -> RolloutItem {
RolloutItem::EventMsg(EventMsg::TokenCount(
codex_protocol::protocol::TokenCountEvent {
info: None,
rate_limits: None,
},
))
}
fn compacted_resume_test_goal_update(thread_id: ThreadId) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(
codex_protocol::protocol::ThreadGoalUpdatedEvent {
thread_id,
turn_id: None,
goal: codex_protocol::protocol::ThreadGoal {
thread_id,
objective: "finish the task".to_string(),
status: codex_protocol::protocol::ThreadGoalStatus::Active,
token_budget: None,
tokens_used: 0,
time_used_seconds: 0,
created_at: 1,
updated_at: 1,
},
},
))
}
fn compacted_resume_test_rollout_contents(items: Vec<RolloutItem>) -> Result<String> {
let mut contents = String::new();
for item in items {
@@ -462,20 +490,24 @@ mod thread_processor_behavior_tests {
let items = read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)?
.expect("compacted resume items");
assert_eq!(items.len(), 6);
assert_eq!(items.len(), 7);
assert!(matches!(items[0], RolloutItem::SessionMeta(_)));
assert!(matches!(
items[1],
RolloutItem::EventMsg(EventMsg::UserMessage(_))
));
assert!(matches!(items[2], RolloutItem::ResponseItem(_)));
assert!(matches!(items[3], RolloutItem::TurnContext(_)));
assert!(matches!(
items[3],
RolloutItem::EventMsg(EventMsg::UserMessage(_))
));
assert!(matches!(items[4], RolloutItem::TurnContext(_)));
assert_eq!(
compacted_resume_test_replacement_text(&items[4]),
compacted_resume_test_replacement_text(&items[5]),
Some("latest checkpoint")
);
assert!(matches!(
&items[5],
&items[6],
RolloutItem::EventMsg(EventMsg::AgentMessage(event))
if event.message == large_tail_message
));
@@ -489,6 +521,67 @@ mod thread_processor_behavior_tests {
Ok(())
}
#[test]
fn compacted_resume_rollout_items_carries_metadata_before_checkpoint() -> Result<()> {
let temp_dir = TempDir::new()?;
let path = temp_dir.path().join("rollout.jsonl");
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
let contents = compacted_resume_test_rollout_contents(vec![
compacted_resume_test_session(thread_id),
compacted_resume_test_user_event("first user event"),
RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")),
compacted_resume_test_token_count(),
compacted_resume_test_goal_update(thread_id),
compacted_resume_test_user_event("middle user event"),
compacted_resume_test_turn_context(),
compacted_resume_test_compaction("latest checkpoint", /*window_id*/ 2),
compacted_resume_test_agent_event("tail"),
])?;
std::fs::write(&path, contents)?;
let byte_len = std::fs::metadata(path.as_path())?.len();
let items = read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)?
.expect("compacted resume items");
assert!(
items
.iter()
.any(|item| matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))))
);
assert!(
items
.iter()
.any(|item| matches!(item, RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(_))))
);
Ok(())
}
#[test]
fn compacted_resume_rollout_items_falls_back_on_rollback_before_checkpoint() -> Result<()> {
let temp_dir = TempDir::new()?;
let path = temp_dir.path().join("rollout.jsonl");
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
let contents = compacted_resume_test_rollout_contents(vec![
compacted_resume_test_session(thread_id),
compacted_resume_test_user_event("first user event"),
RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")),
compacted_resume_test_user_event("middle user event"),
compacted_resume_test_turn_context(),
RolloutItem::EventMsg(EventMsg::ThreadRolledBack(
codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 },
)),
compacted_resume_test_compaction("rolled back checkpoint", /*window_id*/ 2),
compacted_resume_test_agent_event("tail"),
])?;
std::fs::write(&path, contents)?;
let byte_len = std::fs::metadata(path.as_path())?.len();
assert!(
read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)?.is_none()
);
Ok(())
}
#[test]
fn compacted_resume_rollout_items_falls_back_on_malformed_tail_line() -> Result<()> {
let temp_dir = TempDir::new()?;
@@ -562,6 +655,12 @@ mod thread_processor_behavior_tests {
let recent = read_recent_not_loaded_rollout_items(path.as_path(), /*page_size*/ 1)?
.expect("recent notLoaded suffix");
assert!(
recent
.items
.iter()
.any(|item| { matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(_))) })
);
assert!(
recent
.items
@@ -588,6 +687,50 @@ mod thread_processor_behavior_tests {
Ok(())
}
#[test]
fn recent_not_loaded_rollout_items_falls_back_without_explicit_turn_starts() -> Result<()> {
let temp_dir = TempDir::new()?;
let path = temp_dir.path().join("rollout.jsonl");
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
let contents = compacted_resume_test_rollout_contents(vec![
compacted_resume_test_session(thread_id),
compacted_resume_test_user_event("first user event"),
compacted_resume_test_agent_event("first agent event"),
])?;
std::fs::write(&path, contents)?;
assert!(read_recent_not_loaded_rollout_items(path.as_path(), /*page_size*/ 1)?.is_none());
Ok(())
}
#[test]
fn recent_not_loaded_rollout_items_falls_back_on_non_persisted_suffix_item() -> Result<()> {
let temp_dir = TempDir::new()?;
let path = temp_dir.path().join("rollout.jsonl");
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
let contents = compacted_resume_test_rollout_contents(vec![
compacted_resume_test_session(thread_id),
RolloutItem::EventMsg(EventMsg::TurnStarted(
codex_protocol::protocol::TurnStartedEvent {
turn_id: "turn-a".to_string(),
trace_id: None,
started_at: None,
model_context_window: None,
collaboration_mode_kind: Default::default(),
},
)),
compacted_resume_test_user_event("first user event"),
RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::CompactionTrigger {
metadata: None,
}),
compacted_resume_test_agent_event("first agent event"),
])?;
std::fs::write(&path, contents)?;
assert!(read_recent_not_loaded_rollout_items(path.as_path(), /*page_size*/ 1)?.is_none());
Ok(())
}
#[test]
fn validate_dynamic_tools_rejects_empty_namespace() {
let tools = vec![dynamic_tool(