agent/codex
1
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 UUIDv7 context window IDs (#28953)

Browse Source 
## Why

The token-budget context currently identifies a context window by its
thread-local sequence number. A UUIDv7 gives the model a stable opaque
identity that remains fixed for a window and rotates when compaction or
`new_context` starts the next one.

## What changed

- Preserve the existing monotonic value as `window_number` and add a
UUIDv7 `window_id` to `CompactedItem`.
- Generate and rotate the UUID with auto-compaction window state,
persist it alongside the number, and reconstruct it on resume and
rollback.
- Accept legacy compacted rollout records where the numeric `window_id`
represented the window number.
- Use the UUID only in token-budget context; existing request headers
and metadata continue using `thread_id:window_number`.

## Testing

- `just test -p codex-protocol compacted_item::tests`
- `just test -p codex-core token_budget`
This commit is contained in:
pakrym-oai
2026-06-18 17:00:49 -07:00
committed by GitHub
Unverified
parent e83b7841b0
commit 5c12034e42
19 changed files with 319 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/src/protocol/thread_history.rs
+1
View File
@@ -3340,6 +3340,7 @@ mod tests {
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: None,
window_number: None,
window_id: None,
}),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
codex-rs/core/src/agent/control_tests.rs
+1
View File
@@ -1105,6 +1105,7 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() {
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(replacement_history),
window_number: None,
window_id: None,
}),
RolloutItem::TurnContext(turn_context.to_turn_context_item()),
codex-rs/core/src/compact.rs
+2 -1
View File
@@ -302,7 +302,7 @@ async fn run_compact_task_inner_impl(
let user_messages = collect_user_messages(history_items);
let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text);
let window_id = sess.advance_auto_compact_window_id().await;
let (window_number, window_id) = sess.advance_auto_compact_window().await;
if matches!(
initial_context_injection,
@@ -319,6 +319,7 @@ async fn run_compact_task_inner_impl(
let compacted_item = CompactedItem {
message: summary_text.clone(),
replacement_history: Some(new_history.clone()),
window_number: Some(window_number),
window_id: Some(window_id),
};
sess.replace_compacted_history(new_history, reference_context_item, compacted_item)
codex-rs/core/src/compact_remote.rs
+2 -1
View File
@@ -258,7 +258,7 @@ async fn run_remote_compact_task_inner_impl(
&responses_metadata,
)
.await?;
let new_window_id = sess.advance_auto_compact_window_id().await;
let (new_window_number, new_window_id) = sess.advance_auto_compact_window().await;
new_history = process_compacted_history(
sess.as_ref(),
turn_context.as_ref(),
@@ -274,6 +274,7 @@ async fn run_remote_compact_task_inner_impl(
let compacted_item = CompactedItem {
message: String::new(),
replacement_history: Some(new_history.clone()),
window_number: Some(new_window_number),
window_id: Some(new_window_id),
};
// Install is the semantic boundary where the compact endpoint's output becomes live
codex-rs/core/src/compact_remote_v2.rs
+2 -1
View File
@@ -292,7 +292,7 @@ async fn run_remote_compact_task_inner_impl(
let (compacted_history, retained_images) =
build_v2_compacted_history(&prompt_input, compaction_output);
analytics_details.retained_image_count = Some(retained_images);
let new_window_id = sess.advance_auto_compact_window_id().await;
let (new_window_number, new_window_id) = sess.advance_auto_compact_window().await;
let new_history = process_compacted_history(
sess.as_ref(),
turn_context.as_ref(),
@@ -308,6 +308,7 @@ async fn run_remote_compact_task_inner_impl(
let compacted_item = CompactedItem {
message: String::new(),
replacement_history: Some(new_history.clone()),
window_number: Some(new_window_number),
window_id: Some(new_window_id),
};
compaction_trace.record_installed(&CompactionCheckpointTracePayload {
codex-rs/core/src/context/token_budget_context.rs
+4 -3
View File
@@ -1,15 +1,16 @@
use super::ContextualUserFragment;
use codex_protocol::ThreadId;
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct TokenBudgetContext {
thread_id: ThreadId,
window_id: u64,
window_id: Uuid,
tokens_left: i64,
}
impl TokenBudgetContext {
pub(crate) fn new(thread_id: ThreadId, window_id: u64, tokens_left: i64) -> Self {
pub(crate) fn new(thread_id: ThreadId, window_id: Uuid, tokens_left: i64) -> Self {
Self {
thread_id,
window_id,
@@ -36,7 +37,7 @@ impl ContextualUserFragment for TokenBudgetContext {
let window_id = self.window_id;
let tokens_left = self.tokens_left;
format!(
"Thread id {thread_id}.\nCurrent context window {window_id}.\nYou have {tokens_left} tokens left in this context window."
"Thread id {thread_id}.\nCurrent context window id {window_id}.\nYou have {tokens_left} tokens left in this context window."
)
}
}
codex-rs/core/src/session/mod.rs
+13 -9
View File
@@ -1352,6 +1352,7 @@ impl Session {
mut history,
previous_turn_settings,
reference_context_item,
window_number,
window_id,
} = self
.reconstruct_history_from_rollout(turn_context, rollout_items)
@@ -1370,7 +1371,8 @@ impl Session {
{
let mut state = self.state.lock().await;
state.replace_history(history, reference_context_item);
state.set_auto_compact_window_id(window_id);
let window_id = window_id.unwrap_or_else(|| state.auto_compact_window_id());
state.restore_auto_compact_window(window_number, window_id);
state.set_previous_turn_settings(previous_turn_settings.clone());
}
let prefix_tokens = if matches!(
@@ -3231,13 +3233,14 @@ impl Session {
pub(crate) async fn current_window_id(&self) -> String {
let state = self.state.lock().await;
let thread_id = self.thread_id;
let window_id = state.auto_compact_window_id();
format!("{thread_id}:{window_id}")
let window_number = state.auto_compact_window_number();
format!("{thread_id}:{window_number}")
}
pub(crate) async fn advance_auto_compact_window_id(&self) -> u64 {
pub(crate) async fn advance_auto_compact_window(&self) -> (u64, String) {
let mut state = self.state.lock().await;
state.advance_auto_compact_window_id()
let (window_number, window_id) = state.advance_auto_compact_window();
(window_number, window_id.to_string())
}
pub(crate) async fn request_new_context_window(&self) {
@@ -3249,11 +3252,11 @@ impl Session {
&self,
turn_context: &TurnContext,
) -> Option<u64> {
let window_id = {
let window = {
let mut state = self.state.lock().await;
state.start_new_context_window_if_requested()
};
let window_id = window_id?;
let (window_number, window_id) = 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;
@@ -3265,7 +3268,8 @@ impl Session {
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(replacement_history),
window_id: Some(window_id),
window_number: Some(window_number),
window_id: Some(window_id.to_string()),
}),
RolloutItem::TurnContext(turn_context_item),
])
@@ -3275,7 +3279,7 @@ impl Session {
state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact);
}
self.recompute_token_usage(turn_context).await;
Some(window_id)
Some(window_number)
}
pub(crate) async fn reference_context_item(&self) -> Option<TurnContextItem> {
codex-rs/core/src/session/rollout_reconstruction.rs
+36 -12
View File
@@ -1,5 +1,6 @@
use super::*;
use crate::context_manager::is_user_turn_boundary;
use uuid::Uuid;
// Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with
// the resume/fork hydration metadata derived from the same replay.
@@ -8,7 +9,14 @@ pub(super) struct RolloutReconstruction {
pub(super) history: Vec<ResponseItem>,
pub(super) previous_turn_settings: Option<PreviousTurnSettings>,
pub(super) reference_context_item: Option<TurnContextItem>,
pub(super) window_id: u64,
pub(super) window_number: u64,
pub(super) window_id: Option<Uuid>,
}
#[derive(Debug, Clone, Copy)]
struct ReconstructedWindow {
number: u64,
id: Option<Uuid>,
}
#[derive(Debug, Default)]
@@ -34,7 +42,7 @@ struct ActiveReplaySegment<'a> {
previous_turn_settings: Option<PreviousTurnSettings>,
reference_context_item: TurnReferenceContextItem,
base_replacement_history: Option<&'a [ResponseItem]>,
window_id: Option<u64>,
window: Option<ReconstructedWindow>,
}
fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool {
@@ -47,7 +55,7 @@ fn finalize_active_segment<'a>(
base_replacement_history: &mut Option<&'a [ResponseItem]>,
previous_turn_settings: &mut Option<PreviousTurnSettings>,
reference_context_item: &mut TurnReferenceContextItem,
window_id: &mut Option<u64>,
window: &mut Option<ReconstructedWindow>,
pending_rollback_turns: &mut usize,
) {
// Thread rollback drops the newest surviving real user-message boundaries. In replay, that
@@ -68,8 +76,8 @@ fn finalize_active_segment<'a>(
*base_replacement_history = Some(segment_base_replacement_history);
}
if window_id.is_none() {
*window_id = active_segment.window_id;
if window.is_none() {
*window = active_segment.window;
}
// `previous_turn_settings` come from the newest surviving user turn that established them.
@@ -104,7 +112,7 @@ impl Session {
let mut base_replacement_history: Option<&[ResponseItem]> = None;
let mut previous_turn_settings = None;
let mut reference_context_item = TurnReferenceContextItem::NeverSet;
let mut window_id = None;
let mut window = None;
// Rollback is "drop the newest N user turns". While scanning in reverse, that becomes
// "skip the next N user-turn segments we finalize".
let mut pending_rollback_turns = 0usize;
@@ -120,8 +128,13 @@ impl Session {
RolloutItem::Compacted(compacted) => {
let active_segment =
active_segment.get_or_insert_with(ActiveReplaySegment::default);
if active_segment.window_id.is_none() {
active_segment.window_id = compacted.window_id;
if active_segment.window.is_none()
&& let Some(window_number) = compacted.window_number
{
active_segment.window = Some(ReconstructedWindow {
number: window_number,
id: compacted.window_id.as_deref().and_then(parse_uuid_v7),
});
}
// Looking backward, compaction clears any older baseline unless a newer
// `TurnContextItem` in this same segment has already re-established it.
@@ -210,7 +223,7 @@ impl Session {
&mut base_replacement_history,
&mut previous_turn_settings,
&mut reference_context_item,
&mut window_id,
&mut window,
&mut pending_rollback_turns,
);
}
@@ -245,12 +258,12 @@ impl Session {
&mut base_replacement_history,
&mut previous_turn_settings,
&mut reference_context_item,
&mut window_id,
&mut window,
&mut pending_rollback_turns,
);
}
let fallback_window_id = u64::try_from(
let fallback_window_number = u64::try_from(
rollout_items
.iter()
.filter(|item| matches!(item, RolloutItem::Compacted(_)))
@@ -326,11 +339,22 @@ impl Session {
reference_context_item
};
let window = window.unwrap_or(ReconstructedWindow {
number: fallback_window_number,
id: None,
});
RolloutReconstruction {
history: history.raw_items().to_vec(),
previous_turn_settings,
reference_context_item,
window_id: window_id.unwrap_or(fallback_window_id),
window_number: window.number,
window_id: window.id,
}
}
}
fn parse_uuid_v7(value: &str) -> Option<Uuid> {
Uuid::parse_str(value)
.ok()
.filter(|uuid| uuid.get_version_num() == 7)
}
codex-rs/core/src/session/rollout_reconstruction_tests.rs
+8
View File
@@ -826,6 +826,7 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
window_id: None,
}),
RolloutItem::EventMsg(EventMsg::ThreadRolledBack(
@@ -883,6 +884,7 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
window_id: None,
}),
];
@@ -909,6 +911,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_
RolloutItem::Compacted(CompactedItem {
message: "legacy summary".to_string(),
replacement_history: None,
window_number: None,
window_id: None,
}),
];
@@ -941,6 +944,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear
RolloutItem::Compacted(CompactedItem {
message: "legacy summary".to_string(),
replacement_history: None,
window_number: None,
window_id: None,
}),
RolloutItem::EventMsg(EventMsg::TurnStarted(
@@ -1035,6 +1039,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
window_id: None,
}),
RolloutItem::TurnContext(previous_context_item),
@@ -1185,6 +1190,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
window_id: None,
}),
];
@@ -1419,6 +1425,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
window_id: None,
}),
];
@@ -1583,6 +1590,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
window_number: None,
window_id: None,
}),
// A newer TurnStarted replaces the incomplete compacted turn without a matching
codex-rs/core/src/session/tests.rs
+27 -7
View File
@@ -168,6 +168,7 @@ use tokio::sync::Semaphore;
use tokio::time::sleep;
use tokio::time::timeout;
use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid;
use codex_protocol::mcp::CallToolResult as McpCallToolResult;
use pretty_assertions::assert_eq;
@@ -1592,7 +1593,13 @@ async fn reconstruct_history_matches_live_compactions() {
.await;
assert_eq!(expected, reconstructed.history);
assert_eq!(2, reconstructed.window_id);
assert_eq!(2, reconstructed.window_number);
assert_eq!(
reconstructed
.window_id
.map(|window_id| window_id.get_version_num()),
Some(7)
);
}
#[tokio::test]
@@ -1622,10 +1629,12 @@ async fn reconstruct_history_uses_replacement_history_verbatim() {
metadata: None,
},
];
let window_id = Uuid::now_v7();
let rollout_items = vec![RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(replacement_history.clone()),
window_id: Some(42),
window_number: Some(42),
window_id: Some(window_id.to_string()),
})];
let reconstructed = session
@@ -1633,7 +1642,8 @@ async fn reconstruct_history_uses_replacement_history_verbatim() {
.await;
assert_eq!(reconstructed.history, replacement_history);
assert_eq!(42, reconstructed.window_id);
assert_eq!(42, reconstructed.window_number);
assert_eq!(Some(window_id), reconstructed.window_id);
}
#[tokio::test]
@@ -3014,6 +3024,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
user_message("turn 1 user"),
user_message("summary after compaction"),
];
let compacted_window_id = Uuid::now_v7();
sess.persist_rollout_items(&[
RolloutItem::EventMsg(EventMsg::TurnStarted(
@@ -3055,7 +3066,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
RolloutItem::Compacted(CompactedItem {
message: "summary after compaction".to_string(),
replacement_history: Some(compacted_history.clone()),
window_id: Some(7),
window_number: Some(7),
window_id: Some(compacted_window_id.to_string()),
}),
RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: compact_turn_id,
@@ -3105,7 +3117,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio
.await;
{
let mut state = sess.state.lock().await;
state.set_auto_compact_window_id(/*window_id*/ 99);
state.restore_auto_compact_window(/*window_number*/ 99, Uuid::now_v7());
}
handlers::thread_rollback(&sess, "sub-1".to_string(), /*num_turns*/ 1).await;
@@ -3114,6 +3126,10 @@ 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
);
assert!(sess.current_window_id().await.ends_with(":7"));
}
@@ -9846,10 +9862,12 @@ 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;
rollout_items.push(RolloutItem::Compacted(CompactedItem {
message: summary1.to_string(),
replacement_history: None,
window_id: None,
window_number: Some(window_number),
window_id: Some(window_id),
}));
let user2 = ResponseItem::Message {
@@ -9889,10 +9907,12 @@ 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;
rollout_items.push(RolloutItem::Compacted(CompactedItem {
message: summary2.to_string(),
replacement_history: None,
window_id: None,
window_number: Some(window_number),
window_id: Some(window_id),
}));
let user3 = ResponseItem::Message {
codex-rs/core/src/state/auto_compact_window.rs
+28 -12
View File
@@ -1,4 +1,5 @@
use codex_protocol::protocol::TokenUsage;
use uuid::Uuid;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct AutoCompactWindowSnapshot {
@@ -13,7 +14,8 @@ enum AutoCompactWindowPrefill {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) struct AutoCompactWindow {
window_id: u64,
window_number: u64,
window_id: Uuid,
new_context_window_requested: bool,
/// Absolute input-token baseline for the current compaction window.
///
@@ -26,7 +28,8 @@ pub(super) struct AutoCompactWindow {
impl AutoCompactWindow {
pub(super) fn new() -> Self {
Self {
window_id: 0,
window_number: 0,
window_id: Uuid::now_v7(),
new_context_window_requested: false,
prefill_input_tokens: None,
}
@@ -36,18 +39,24 @@ impl AutoCompactWindow {
self.prefill_input_tokens = None;
}
pub(super) fn window_id(&self) -> u64 {
pub(super) fn window_number(&self) -> u64 {
self.window_number
}
pub(super) fn window_id(&self) -> Uuid {
self.window_id
}
pub(super) fn set_window_id(&mut self, window_id: u64) {
pub(super) fn restore(&mut self, window_number: u64, window_id: Uuid) {
self.window_number = window_number;
self.window_id = window_id;
}
pub(super) fn advance_window_id(&mut self) -> u64 {
self.window_id = self.window_id.saturating_add(1);
pub(super) fn advance(&mut self) -> (u64, Uuid) {
self.window_number = self.window_number.saturating_add(1);
self.window_id = Uuid::now_v7();
self.new_context_window_requested = false;
self.window_id
(self.window_number, self.window_id)
}
pub(super) fn request_new_context_window(&mut self) {
@@ -108,15 +117,22 @@ mod tests {
fn tracks_prefill_and_window_boundaries() {
let mut window = AutoCompactWindow::new();
assert_eq!(window.window_id(), 0);
window.set_window_id(/*window_id*/ 3);
assert_eq!(window.window_id(), 3);
assert_eq!(window.window_number(), 0);
assert_eq!(window.window_id().get_version_num(), 7);
let restored_window_id = Uuid::now_v7();
window.restore(/*window_number*/ 3, restored_window_id);
assert_eq!(window.window_number(), 3);
assert_eq!(window.window_id(), restored_window_id);
window.request_new_context_window();
assert!(window.take_new_context_window_request());
assert!(!window.take_new_context_window_request());
window.request_new_context_window();
assert_eq!(window.advance_window_id(), 4);
assert_eq!(window.window_id(), 4);
let (window_number, window_id) = window.advance();
assert_eq!(window_number, 4);
assert_eq!(window.window_number(), 4);
assert_eq!(window.window_id(), window_id);
assert_eq!(window_id.get_version_num(), 7);
assert_ne!(window_id, restored_window_id);
assert!(!window.take_new_context_window_request());
assert_eq!(
codex-rs/core/src/state/session.rs
+13 -8
View File
@@ -6,6 +6,7 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles;
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use uuid::Uuid;
use super::AdditionalContextStore;
use super::auto_compact_window::AutoCompactWindow;
@@ -147,30 +148,34 @@ impl SessionState {
self.auto_compact_window.snapshot()
}
pub(crate) fn auto_compact_window_id(&self) -> u64 {
pub(crate) fn auto_compact_window_number(&self) -> u64 {
self.auto_compact_window.window_number()
}
pub(crate) fn auto_compact_window_id(&self) -> Uuid {
self.auto_compact_window.window_id()
}
pub(crate) fn set_auto_compact_window_id(&mut self, window_id: u64) {
self.auto_compact_window.set_window_id(window_id);
pub(crate) fn restore_auto_compact_window(&mut self, window_number: u64, window_id: Uuid) {
self.auto_compact_window.restore(window_number, window_id);
}
pub(crate) fn advance_auto_compact_window_id(&mut self) -> u64 {
self.auto_compact_window.advance_window_id()
pub(crate) fn advance_auto_compact_window(&mut self) -> (u64, Uuid) {
self.auto_compact_window.advance()
}
pub(crate) fn request_new_context_window(&mut self) {
self.auto_compact_window.request_new_context_window();
}
pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option<u64> {
pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option<(u64, Uuid)> {
if !self.auto_compact_window.take_new_context_window_request() {
return None;
}
let window_id = self.auto_compact_window.advance_window_id();
let window = self.auto_compact_window.advance();
self.auto_compact_window.clear_prefill();
Some(window_id)
Some(window)
}
pub(crate) fn token_info(&self) -> Option<TokenUsageInfo> {
codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap
+1 -1
View File
@@ -8,7 +8,7 @@ Scenario: New context window tool installs fresh full context before the next fo
00:message/developer[3]:
[01] <PERMISSIONS_INSTRUCTIONS>
[02] <SKILLS_INSTRUCTIONS>
[03] <token_budget>\nThread id <THREAD_ID>.\nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n</token_budget>
[03] <token_budget>\nThread id <THREAD_ID>.\nCurrent context window id <UUID>.\nYou have 121600 tokens left in this context window.\n</token_budget>
01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
02:function_call/update_plan
03:function_call_output:Plan updated
codex-rs/core/tests/suite/token_budget.rs
+78 -30
View File
@@ -4,6 +4,7 @@ use codex_model_provider_info::built_in_model_providers;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use core_test_support::PathBufExt;
use core_test_support::assert_regex_match;
use core_test_support::context_snapshot;
use core_test_support::context_snapshot::ContextSnapshotOptions;
use core_test_support::responses::ResponsesRequest;
@@ -80,17 +81,17 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()>
assert_eq!(requests.len(), 2);
let thread_id = test.session_configured.thread_id;
let expected = vec![format!(
"<token_budget>\nThread id {thread_id}.\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
)];
assert_eq!(
token_budget_texts(&requests[0]),
expected,
"initial full context should report context window 0"
let initial_token_budget = token_budget_texts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id [0-9a-f-]{{36}}\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n</token_budget>$"
),
&initial_token_budget[0],
);
assert_eq!(
token_budget_texts(&requests[1]),
expected,
initial_token_budget,
"steady-state context update should not advance the context window"
);
@@ -144,9 +145,15 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R
assert_eq!(requests.len(), 5);
let thread_id = test.session_configured.thread_id;
let full_context = format!(
"<token_budget>\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
let full_context = token_budget_texts(&requests[0]);
assert_eq!(full_context.len(), 1);
assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id [0-9a-f-]{{36}}\.\nYou have 9500 tokens left in this context window\.\n</token_budget>$"
),
&full_context[0],
);
let full_context = full_context[0].clone();
let threshold_25 =
"<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
.to_string();
@@ -233,16 +240,18 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu
);
let thread_id = test.session_configured.thread_id;
let full_context = format!(
"<token_budget>\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n</token_budget>"
);
let remaining_context =
"<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
.to_string();
assert_eq!(
token_budget_texts(&requests[1]),
vec![full_context, remaining_context.clone()]
let token_budgets = token_budget_texts(&requests[1]);
assert_eq!(token_budgets.len(), 2);
assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id [0-9a-f-]{{36}}\.\nYou have 9500 tokens left in this context window\.\n</token_budget>$"
),
&token_budgets[0],
);
assert_eq!(token_budgets[1], remaining_context);
assert_eq!(
requests[2].function_call_output_content_and_success(call_id),
Some((Some(remaining_context), None))
@@ -362,13 +371,31 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
assert_eq!(requests.len(), 3);
let thread_id = test.session_configured.thread_id;
assert_eq!(
token_budget_texts(&requests[2]),
vec![format!(
"<token_budget>\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
)],
"post-compaction full context should report context window 1"
);
let initial_token_budget = token_budget_texts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let initial_window_id = assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n</token_budget>$"
),
&initial_token_budget[0],
)
.get(1)
.expect("window id capture")
.as_str()
.to_string();
let post_compaction_token_budget = token_budget_texts(&requests[2]);
assert_eq!(post_compaction_token_budget.len(), 1);
let post_compaction_window_id = assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n</token_budget>$"
),
&post_compaction_token_budget[0],
)
.get(1)
.expect("window id capture")
.as_str()
.to_string();
assert_ne!(post_compaction_window_id, initial_window_id);
Ok(())
}
@@ -429,12 +456,31 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
"new_context should be exposed when token budget is enabled"
);
let thread_id = test.session_configured.thread_id;
assert_eq!(
token_budget_texts(&requests[2]),
vec![format!(
"<token_budget>\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
)]
);
let initial_token_budget = token_budget_texts(&requests[0]);
assert_eq!(initial_token_budget.len(), 1);
let initial_window_id = assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n</token_budget>$"
),
&initial_token_budget[0],
)
.get(1)
.expect("window id capture")
.as_str()
.to_string();
let new_window_token_budget = token_budget_texts(&requests[2]);
assert_eq!(new_window_token_budget.len(), 1);
let window_id = assert_regex_match(
&format!(
r"^<token_budget>\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n</token_budget>$"
),
&new_window_token_budget[0],
)
.get(1)
.expect("window id capture")
.as_str()
.to_string();
assert_ne!(window_id, initial_window_id);
assert!(
!requests[2].body_contains_text("request new context window"),
"new_context should drop the prior window history before continuing the turn"
@@ -448,7 +494,9 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
&[("Final Follow-Up Request", &requests[2])],
&ContextSnapshotOptions::default(),
);
let snapshot = snapshot.replace(&thread_id.to_string(), "<THREAD_ID>");
let snapshot = snapshot
.replace(&thread_id.to_string(), "<THREAD_ID>")
.replace(&window_id, "<UUID>");
insta::assert_snapshot!(
"token_budget_new_context_window_tool_full_context",
snapshot
codex-rs/protocol/src/compacted_item.rs
+94
View File
@@ -0,0 +1,94 @@
use crate::models::ResponseItem;
use crate::protocol::CompactedItem;
use serde::Deserialize;
// Before `window_number` was introduced, the numeric window number was serialized as
// `window_id`. Accept that shape so existing rollouts remain resumable.
impl<'de> Deserialize<'de> for CompactedItem {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let serialized = SerializedCompactedItem::deserialize(deserializer)?;
let mut window_number = serialized.window_number;
let window_id = match serialized.window_id {
Some(SerializedWindowId::Id(window_id)) => Some(window_id),
Some(SerializedWindowId::LegacyWindowNumber(legacy_window_number)) => {
window_number.get_or_insert(legacy_window_number);
None
}
None => None,
};
Ok(Self {
message: serialized.message,
replacement_history: serialized.replacement_history,
window_number,
window_id,
})
}
}
#[derive(Deserialize)]
struct SerializedCompactedItem {
message: String,
#[serde(default)]
replacement_history: Option<Vec<ResponseItem>>,
#[serde(default)]
window_number: Option<u64>,
#[serde(default)]
window_id: Option<SerializedWindowId>,
}
#[derive(Deserialize)]
#[serde(untagged)]
enum SerializedWindowId {
Id(String),
LegacyWindowNumber(u64),
}
#[cfg(test)]
mod tests {
use super::*;
use anyhow::Result;
use pretty_assertions::assert_eq;
use serde_json::json;
#[test]
fn serializes_window_number_and_id() -> Result<()> {
let item = CompactedItem {
message: "summary".to_string(),
replacement_history: None,
window_number: Some(3),
window_id: Some("019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001".to_string()),
};
assert_eq!(
serde_json::to_value(item)?,
json!({
"message": "summary",
"window_number": 3,
"window_id": "019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001",
})
);
Ok(())
}
#[test]
fn migrates_legacy_numeric_window_id() -> Result<()> {
let item = serde_json::from_value::<CompactedItem>(json!({
"message": "summary",
"window_id": 3,
}))?;
assert_eq!(
item,
CompactedItem {
message: "summary".to_string(),
replacement_history: None,
window_number: Some(3),
window_id: None,
}
);
Ok(())
}
}
codex-rs/protocol/src/lib.rs
+1
View File
@@ -10,6 +10,7 @@ pub use thread_id::ThreadId;
pub use tool_name::ToolName;
pub mod approvals;
pub mod capabilities;
mod compacted_item;
pub mod config_types;
pub mod dynamic_tools;
pub mod error;
codex-rs/protocol/src/protocol.rs
+6 -2
View File
@@ -2971,13 +2971,17 @@ pub enum RolloutItem {
EventMsg(EventMsg),
}
#[derive(Serialize, Deserialize, Clone, Debug, JsonSchema, TS)]
#[derive(Serialize, Clone, Debug, PartialEq, JsonSchema, TS)]
pub struct CompactedItem {
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub replacement_history: Option<Vec<ResponseItem>>,
/// Monotonic position of this context window within the thread.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub window_id: Option<u64>,
pub window_number: Option<u64>,
/// UUIDv7 identity of this context window.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub window_id: Option<String>,
}
impl From<CompactedItem> for ResponseItem {
codex-rs/rollout/src/metadata_tests.rs
+1
View File
@@ -154,6 +154,7 @@ fn builder_from_items_falls_back_to_filename() {
let items = vec![RolloutItem::Compacted(CompactedItem {
message: "noop".to_string(),
replacement_history: None,
window_number: None,
window_id: None,
})];
codex-rs/thread-store/src/thread_metadata_sync.rs
+1
View File
@@ -483,6 +483,7 @@ mod tests {
let item = RolloutItem::Compacted(CompactedItem {
message: "compacted".to_string(),
replacement_history: None,
window_number: None,
window_id: None,
});