openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-04 15:08:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: extract context window token status

Browse Source
This commit is contained in:
Michael Bolin
2026-06-23 13:12:22 -07:00
parent 00dc5ea397
commit ed6ce04ea2
3 changed files with 86 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
codex-rs/core/src/session/context_window.rs Normal file
View File

@@ -0,0 +1,76 @@
use super::session::Session;
use super::turn_context::TurnContext;
use codex_protocol::config_types::AutoCompactTokenLimitScope;
#[derive(Debug)]
pub(crate) struct ContextWindowTokenStatus {
// Full active context usage, independent of the configured auto-compact scope.
pub(crate) active_context_tokens: i64,
// Usage counted against `model_auto_compact_token_limit` for the current scope.
pub(crate) auto_compact_scope_tokens: i64,
pub(crate) auto_compact_scope_limit: i64,
pub(crate) full_context_window_limit: Option<i64>,
pub(crate) auto_compact_window_prefill_tokens: Option<i64>,
pub(crate) full_context_window_limit_reached: bool,
pub(crate) token_limit_reached: bool,
}
impl ContextWindowTokenStatus {
pub(crate) fn tokens_until_compaction(&self) -> i64 {
let full_context_remaining = self.full_context_window_limit.map_or(i64::MAX, |limit| {
limit.saturating_sub(self.active_context_tokens)
});
self.auto_compact_scope_limit
.saturating_sub(self.auto_compact_scope_tokens)
.min(full_context_remaining)
.max(0)
}
}
pub(crate) async fn context_window_token_status(
sess: &Session,
turn_context: &TurnContext,
) -> ContextWindowTokenStatus {
let active_context_tokens = sess.get_total_token_usage().await;
let mut auto_compact_window_prefill_tokens = None;
let (auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit) =
match turn_context.config.model_auto_compact_token_limit_scope {
AutoCompactTokenLimitScope::Total => (
active_context_tokens,
turn_context
.model_info
.auto_compact_token_limit()
.unwrap_or(i64::MAX),
None,
),
AutoCompactTokenLimitScope::BodyAfterPrefix => {
let window = sess.auto_compact_window_snapshot().await;
auto_compact_window_prefill_tokens = window.prefill_input_tokens;
let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens);
(
active_context_tokens.saturating_sub(baseline),
turn_context
.config
.model_auto_compact_token_limit
.or_else(|| turn_context.model_info.auto_compact_token_limit())
.unwrap_or(i64::MAX),
turn_context.model_context_window(),
)
}
};
let full_context_window_limit_reached =
full_context_window_limit.is_some_and(|full_context_window_limit| {
active_context_tokens >= full_context_window_limit
});
let token_limit_reached =
auto_compact_scope_tokens >= auto_compact_scope_limit || full_context_window_limit_reached;
ContextWindowTokenStatus {
active_context_tokens,
auto_compact_scope_tokens,
auto_compact_scope_limit,
full_context_window_limit,
auto_compact_window_prefill_tokens,
full_context_window_limit_reached,
token_limit_reached,
}
}

1
codex-rs/core/src/session/mod.rs
View File

@@ -210,6 +210,7 @@ use codex_protocol::exec_output::StreamOutput;
mod code_mode_warning;
mod config_lock;
pub(crate) mod context_window;
mod handlers;
mod inject;
mod input_queue;

85
codex-rs/core/src/session/turn.rs
View File

@@ -298,8 +298,11 @@ pub(crate) async fn run_turn(
let (has_pending_input, token_status, estimated_token_count) = async {
let has_pending_input =
sess.input_queue.has_pending_input(&sess.active_turn).await;
let token_status =
auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await;
let token_status = super::context_window::context_window_token_status(
sess.as_ref(),
turn_context.as_ref(),
)
.await;
let estimated_token_count =
sess.get_estimated_token_count(turn_context.as_ref()).await;
(has_pending_input, token_status, estimated_token_count)
@@ -326,17 +329,7 @@ pub(crate) async fn run_turn(
"post sampling token usage"
);
let tokens_after_sampling = token_status.active_context_tokens;
let full_context_remaining = token_status
.full_context_window_limit
.map_or(i64::MAX, |limit| {
limit.saturating_sub(tokens_after_sampling)
});
let tokens_until_compaction = token_status
.auto_compact_scope_limit
.saturating_sub(token_status.auto_compact_scope_tokens)
.min(full_context_remaining)
.max(0);
let tokens_until_compaction = token_status.tokens_until_compaction();
super::token_budget::maybe_record(
sess.as_ref(),
turn_context.as_ref(),
@@ -806,68 +799,6 @@ async fn track_turn_resolved_config_analytics(
});
}
#[derive(Debug)]
struct AutoCompactTokenStatus {
// Full active context usage, independent of the configured auto-compact scope.
active_context_tokens: i64,
// Usage counted against `model_auto_compact_token_limit` for the current scope.
auto_compact_scope_tokens: i64,
auto_compact_scope_limit: i64,
full_context_window_limit: Option<i64>,
auto_compact_window_prefill_tokens: Option<i64>,
full_context_window_limit_reached: bool,
token_limit_reached: bool,
}
async fn auto_compact_token_status(
sess: &Session,
turn_context: &TurnContext,
) -> AutoCompactTokenStatus {
let active_context_tokens = sess.get_total_token_usage().await;
let mut auto_compact_window_prefill_tokens = None;
let (auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit) =
match turn_context.config.model_auto_compact_token_limit_scope {
AutoCompactTokenLimitScope::Total => (
active_context_tokens,
turn_context
.model_info
.auto_compact_token_limit()
.unwrap_or(i64::MAX),
None,
),
AutoCompactTokenLimitScope::BodyAfterPrefix => {
let window = sess.auto_compact_window_snapshot().await;
auto_compact_window_prefill_tokens = window.prefill_input_tokens;
let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens);
(
active_context_tokens.saturating_sub(baseline),
turn_context
.config
.model_auto_compact_token_limit
.or_else(|| turn_context.model_info.auto_compact_token_limit())
.unwrap_or(i64::MAX),
turn_context.model_context_window(),
)
}
};
let full_context_window_limit_reached =
full_context_window_limit.is_some_and(|full_context_window_limit| {
active_context_tokens >= full_context_window_limit
});
let token_limit_reached =
auto_compact_scope_tokens >= auto_compact_scope_limit || full_context_window_limit_reached;
AutoCompactTokenStatus {
active_context_tokens,
auto_compact_scope_tokens,
auto_compact_scope_limit,
full_context_window_limit,
auto_compact_window_prefill_tokens,
full_context_window_limit_reached,
token_limit_reached,
}
}
#[instrument(level = "trace", skip_all)]
async fn run_pre_sampling_compact(
sess: &Arc<Session>,
@@ -883,7 +814,9 @@ async fn run_pre_sampling_compact(
}
maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?;
let token_status = auto_compact_token_status(sess.as_ref(), turn_context.as_ref()).await;
let token_status =
super::context_window::context_window_token_status(sess.as_ref(), turn_context.as_ref())
.await;
// Compact if the configured auto-compaction budget or usable context window is exhausted.
if token_status.token_limit_reached {
run_auto_compact(