Files
codex/codex-rs/core/src/session/context_window.rs
T
rka-oai 768330dd6c Add a fallback phase before automatic context rollover (#33255)
## Why

Token-budget sessions need a final opportunity to preserve important state before the current context is compacted.

## What changed

- When `auto_compact_fallback_prompt` is configured, reserve `auto_compact_fallback_buffer_tokens` beyond the base auto-compaction limit.
- At the base limit, inject the fallback prompt once as a developer message while keeping the normal tool surface available. Continue reporting zero base-window tokens during this reserve.
- Roll over when the reserve or full model context is exhausted, and skip the fallback when a new context was already requested.

## Testing

Added coverage for using the fallback reserve, rolling over after it is exhausted, and bypassing the fallback for an explicit `new_context` request.

GitOrigin-RevId: 32983ecadc5cb4927e15ecd7b1e896cb52a9ad3b
2026-07-15 05:45:43 +00:00

92 lines
3.8 KiB
Rust

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: Option<i64>,
pub(crate) full_context_window_limit: Option<i64>,
pub(crate) base_window_tokens_remaining: 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,
}
fn tokens_remaining(limit: Option<i64>, used: i64) -> Option<i64> {
limit.map(|limit| limit.saturating_sub(used).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;
// Count either the full active context or only the tokens added after the initial prefix.
let (auto_compact_scope_tokens, auto_compact_scope_limit, auto_compact_window_prefill_tokens) =
match turn_context.config.model_auto_compact_token_limit_scope {
AutoCompactTokenLimitScope::Total => (
active_context_tokens,
turn_context.model_info.auto_compact_token_limit(),
None,
),
AutoCompactTokenLimitScope::BodyAfterPrefix => {
let window = sess.auto_compact_window_snapshot().await;
let baseline = window.prefill_input_tokens.unwrap_or(active_context_tokens);
let scope_limit = turn_context
.config
.model_auto_compact_token_limit
.or_else(|| turn_context.model_info.auto_compact_token_limit());
(
active_context_tokens.saturating_sub(baseline),
scope_limit,
window.prefill_input_tokens,
)
}
};
// The model's full context window is a hard cap, independent of the auto-compaction scope.
let full_context_window_limit = turn_context.model_context_window();
// Report remaining tokens against the base (unbuffered) window, capped by the full context.
let base_window_tokens_remaining = [
tokens_remaining(auto_compact_scope_limit, auto_compact_scope_tokens),
tokens_remaining(full_context_window_limit, active_context_tokens),
]
.into_iter()
.flatten()
.min();
// Only reserve the fallback buffer when there is a fallback prompt to use it.
let auto_compact_fallback_buffer_tokens = turn_context
.config
.token_budget
.as_ref()
.map_or(0, crate::config::TokenBudgetConfig::fallback_buffer_tokens);
let buffered_auto_compact_limit = auto_compact_scope_limit
.map(|limit| limit.saturating_add(auto_compact_fallback_buffer_tokens));
// Force compaction once the buffered window or the model's full context window is reached.
let full_context_window_limit_reached =
full_context_window_limit.is_some_and(|limit| active_context_tokens >= limit);
let token_limit_reached = buffered_auto_compact_limit
.is_some_and(|limit| auto_compact_scope_tokens >= limit)
|| full_context_window_limit_reached;
ContextWindowTokenStatus {
active_context_tokens,
auto_compact_scope_tokens,
auto_compact_scope_limit,
full_context_window_limit,
base_window_tokens_remaining,
auto_compact_window_prefill_tokens,
full_context_window_limit_reached,
token_limit_reached,
}
}