agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

core: add context window lineage IDs (#29256)

Browse Source 
## Why

The rendered `<token_budget>` fragment identifies the thread and current
context window, but it does not expose enough lineage to identify the
first window in the thread or the immediately preceding window. Those
IDs also need to remain stable across compaction, resume, and rollback.

## What changed

- Track first, previous, and current UUIDv7 context-window IDs in
auto-compaction state.
- Render `thread_id`, `first_window_id`, `previous_window_id`, and the
current window ID in the full `<token_budget>` fragment.
- Persist the first and previous window IDs in compacted rollout
checkpoints and restore them during rollout reconstruction.
- Preserve compatibility with older compacted records that do not
contain the new optional fields.
- Update focused state, rendering, reconstruction, rollback, and
serialization coverage.

## Validation

- `just test -p codex-core token_budget`
- `just test -p codex-protocol compacted_item::tests`
- `just test -p codex-core tracks_prefill_and_window_boundaries`
- `just test -p codex-core
reconstruct_history_uses_replacement_history_verbatim`
- `just test -p codex-core
thread_rollback_restores_cleared_reference_context_item_after_compaction`
This commit is contained in:
pakrym-oai
2026-06-20 13:15:49 -07:00
committed by GitHub
Unverified
parent d667082322
commit d1209bddfc
19 changed files with 277 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/mod.rs
+24 -10
View File
@@ -311,6 +311,7 @@ use crate::session_startup_prewarm::SessionStartupPrewarmHandle;
use crate::shell;
#[cfg(test)]
use crate::skills::SkillLoadOutcome;
use crate::state::AutoCompactWindowIds;
use crate::state::AutoCompactWindowSnapshot;
use crate::state::PendingRequestPermissions;
use crate::state::SessionServices;
@@ -1364,6 +1365,8 @@ impl Session {
previous_turn_settings,
reference_context_item,
window_number,
first_window_id,
previous_window_id,
window_id,
} = self
.reconstruct_history_from_rollout(turn_context, rollout_items)
@@ -1382,8 +1385,16 @@ impl Session {
{
let mut state = self.state.lock().await;
state.replace_history(history, reference_context_item);
let window_id = window_id.unwrap_or_else(|| state.auto_compact_window_id());
state.restore_auto_compact_window(window_number, window_id);
let fallback_ids = state.auto_compact_window_ids();
let window_id = window_id.unwrap_or(fallback_ids.window_id);
state.restore_auto_compact_window(
window_number,
AutoCompactWindowIds {
first_window_id: first_window_id.unwrap_or(window_id),
previous_window_id,
window_id,
},
);
state.set_previous_turn_settings(previous_turn_settings.clone());
}
let prefix_tokens = if matches!(
@@ -3011,7 +3022,7 @@ impl Session {
collaboration_mode,
base_instructions,
session_source,
auto_compact_window_id,
auto_compact_window_ids,
) = {
let state = self.state.lock().await;
(
@@ -3020,7 +3031,7 @@ impl Session {
state.session_configuration.collaboration_mode.clone(),
state.session_configuration.base_instructions.clone(),
state.session_configuration.session_source.clone(),
state.auto_compact_window_id(),
state.auto_compact_window_ids(),
)
};
if let Some(model_switch_message) =
@@ -3210,7 +3221,9 @@ impl Session {
developer_sections.push(
crate::context::TokenBudgetContext::new(
self.thread_id(),
auto_compact_window_id,
auto_compact_window_ids.first_window_id,
auto_compact_window_ids.previous_window_id,
auto_compact_window_ids.window_id,
model_context_window,
)
.render(),
@@ -3309,10 +3322,9 @@ impl Session {
format!("{thread_id}:{window_number}")
}
pub(crate) async fn advance_auto_compact_window(&self) -> (u64, String) {
pub(crate) async fn advance_auto_compact_window(&self) -> (u64, AutoCompactWindowIds) {
let mut state = self.state.lock().await;
let (window_number, window_id) = state.advance_auto_compact_window();
(window_number, window_id.to_string())
state.advance_auto_compact_window()
}
pub(crate) async fn request_new_context_window(&self) {
@@ -3328,7 +3340,7 @@ impl Session {
let mut state = self.state.lock().await;
state.start_new_context_window_if_requested()
};
let (window_number, window_id) = window?;
let (window_number, window_ids) = window?;
let context_items = self.build_initial_context(turn_context).await;
let turn_context_item = turn_context.to_turn_context_item();
let replacement_history = context_items;
@@ -3341,7 +3353,9 @@ impl Session {
message: String::new(),
replacement_history: Some(replacement_history),
window_number: Some(window_number),
window_id: Some(window_id.to_string()),
first_window_id: Some(window_ids.first_window_id.to_string()),
previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()),
window_id: Some(window_ids.window_id.to_string()),
}),
RolloutItem::TurnContext(turn_context_item),
])
codex-rs/core/src/session/rollout_reconstruction.rs
+13
View File
@@ -10,12 +10,16 @@ pub(super) struct RolloutReconstruction {
pub(super) previous_turn_settings: Option<PreviousTurnSettings>,
pub(super) reference_context_item: Option<TurnContextItem>,
pub(super) window_number: u64,
pub(super) first_window_id: Option<Uuid>,
pub(super) previous_window_id: Option<Uuid>,
pub(super) window_id: Option<Uuid>,
}
#[derive(Debug, Clone, Copy)]
struct ReconstructedWindow {
number: u64,
first_id: Option<Uuid>,
previous_id: Option<Uuid>,
id: Option<Uuid>,
}
@@ -133,6 +137,11 @@ impl Session {
{
active_segment.window = Some(ReconstructedWindow {
number: window_number,
first_id: compacted.first_window_id.as_deref().and_then(parse_uuid_v7),
previous_id: compacted
.previous_window_id
.as_deref()
.and_then(parse_uuid_v7),
id: compacted.window_id.as_deref().and_then(parse_uuid_v7),
});
}
@@ -341,6 +350,8 @@ impl Session {
let window = window.unwrap_or(ReconstructedWindow {
number: fallback_window_number,
first_id: None,
previous_id: None,
id: None,
});
RolloutReconstruction {
@@ -348,6 +359,8 @@ impl Session {
previous_turn_settings,
reference_context_item,
window_number: window.number,
first_window_id: window.first_id,
previous_window_id: window.previous_id,
window_id: window.id,
}
}
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+16
View File
@@ -829,6 +829,8 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
RolloutItem::EventMsg(EventMsg::ThreadRolledBack(
@@ -887,6 +889,8 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
];
@@ -914,6 +918,8 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_
message: "legacy summary".to_string(),
replacement_history: None,
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
];
@@ -947,6 +953,8 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear
message: "legacy summary".to_string(),
replacement_history: None,
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
RolloutItem::EventMsg(EventMsg::TurnStarted(
@@ -1043,6 +1051,8 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
RolloutItem::TurnContext(previous_context_item),
@@ -1196,6 +1206,8 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
];
@@ -1433,6 +1445,8 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
];
@@ -1599,6 +1613,8 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
first_window_id: None,
previous_window_id: None,
window_id: None,
}),
// A newer TurnStarted replaces the incomplete compacted turn without a matching
codex-rs/core/src/session/tests.rs
+32 -7
View File
@@ -1653,11 +1653,15 @@ async fn reconstruct_history_uses_replacement_history_verbatim() {
metadata: None,
},
];
let first_window_id = Uuid::now_v7();
let previous_window_id = Uuid::now_v7();
let window_id = Uuid::now_v7();
let rollout_items = vec![RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(replacement_history.clone()),
window_number: Some(42),
first_window_id: Some(first_window_id.to_string()),
previous_window_id: Some(previous_window_id.to_string()),
window_id: Some(window_id.to_string()),
})];
@@ -1667,6 +1671,8 @@ async fn reconstruct_history_uses_replacement_history_verbatim() {
assert_eq!(reconstructed.history, replacement_history);
assert_eq!(42, reconstructed.window_number);
assert_eq!(Some(first_window_id), reconstructed.first_window_id);
assert_eq!(Some(previous_window_id), reconstructed.previous_window_id);
assert_eq!(Some(window_id), reconstructed.window_id);
}
@@ -3059,6 +3065,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
user_message("turn 1 user"),
user_message("summary after compaction"),
];
let first_window_id = Uuid::now_v7();
let previous_window_id = Uuid::now_v7();
let compacted_window_id = Uuid::now_v7();
sess.persist_rollout_items(&[
@@ -3102,6 +3110,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
message: "summary after compaction".to_string(),
replacement_history: Some(compacted_history.clone()),
window_number: Some(7),
first_window_id: Some(first_window_id.to_string()),
previous_window_id: Some(previous_window_id.to_string()),
window_id: Some(compacted_window_id.to_string()),
}),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
@@ -3152,7 +3162,14 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
.await;
{
let mut state = sess.state.lock().await;
state.restore_auto_compact_window(/*window_number*/ 99, Uuid::now_v7());
state.restore_auto_compact_window(
/*window_number*/ 99,
AutoCompactWindowIds {
first_window_id: Uuid::now_v7(),
previous_window_id: Some(Uuid::now_v7()),
window_id: Uuid::now_v7(),
},
);
}
handlers::thread_rollback(&sess, "sub-1".to_string(), /*num_turns*/ 1).await;
@@ -3162,8 +3179,12 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
assert_eq!(sess.clone_history().await.raw_items(), compacted_history);
assert!(sess.reference_context_item().await.is_none());
assert_eq!(
sess.state.lock().await.auto_compact_window_id(),
compacted_window_id
sess.state.lock().await.auto_compact_window_ids(),
AutoCompactWindowIds {
first_window_id,
previous_window_id: Some(previous_window_id),
window_id: compacted_window_id,
}
);
assert!(sess.current_window_id().await.ends_with(":7"));
}
@@ -9910,12 +9931,14 @@ async fn sample_rollout(
let user_messages1 = collect_user_messages(&snapshot1);
let rebuilt1 = compact::build_compacted_history(Vec::new(), &user_messages1, summary1);
live_history.replace(rebuilt1);
let (window_number, window_id) = session.advance_auto_compact_window().await;
let (window_number, window_ids) = session.advance_auto_compact_window().await;
rollout_items.push(RolloutItem::Compacted(CompactedItem {
message: summary1.to_string(),
replacement_history: None,
window_number: Some(window_number),
window_id: Some(window_id),
first_window_id: Some(window_ids.first_window_id.to_string()),
previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()),
window_id: Some(window_ids.window_id.to_string()),
}));
let user2 = ResponseItem::Message {
@@ -9955,12 +9978,14 @@ async fn sample_rollout(
let user_messages2 = collect_user_messages(&snapshot2);
let rebuilt2 = compact::build_compacted_history(Vec::new(), &user_messages2, summary2);
live_history.replace(rebuilt2);
let (window_number, window_id) = session.advance_auto_compact_window().await;
let (window_number, window_ids) = session.advance_auto_compact_window().await;
rollout_items.push(RolloutItem::Compacted(CompactedItem {
message: summary2.to_string(),
replacement_history: None,
window_number: Some(window_number),
window_id: Some(window_id),
first_window_id: Some(window_ids.first_window_id.to_string()),
previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()),
window_id: Some(window_ids.window_id.to_string()),
}));
let user3 = ResponseItem::Message {