mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
refactor: extract context window token status (#29664)
## Why This PR keeps the mechanical helper extraction separate from the behavior change in #29665. The follow-up needs the token-window accounting from `turn.rs` in another call path, but reviewing that is much easier when the helper extraction is separate from the semantic change. ## What - Adds `session/context_window.rs` with `ContextWindowTokenStatus`. - Moves the existing auto-compaction token-status calculation out of `session/turn.rs`. - Replaces the duplicated inline remaining-token calculation in `turn.rs` with `tokens_until_compaction()`. This PR is intended to be behavior-preserving. The `get_context_remaining` behavior change is stacked separately in #29665. ## Testing - `just test -p codex-core auto_compact_body_after_prefix` --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/29664). * #29665 * __->__ #29664
This commit is contained in:
committed by
GitHub
Unverified
parent
829f5b6b59
commit
4dde907d27
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -211,6 +211,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;
|
||||
|
||||
@@ -293,8 +293,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)
|
||||
@@ -321,17 +324,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(),
|
||||
@@ -801,68 +794,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>,
|
||||
@@ -878,7 +809,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 {
|
||||
let step_context = sess.capture_step_context(Arc::clone(turn_context)).await;
|
||||
|
||||
Reference in New Issue
Block a user