[codex] simplify token budget context (#29295)

## Why

The token-budget feature currently adds remaining-token messages
whenever usage crosses the 25%, 50%, and 75% thresholds. Those periodic
inserts create prompt churn without requiring action, while the
near-compaction reminder and explicit `get_context_remaining` tool
already cover actionable and on-demand budget information.

The context-window lineage block is also easier to scan as plain labeled
text than as a `<token_budget>`-wrapped fragment.

## What changed

- Stop recording automatic remaining-token messages at percentage
thresholds.
- Render context-window lineage in `First`, `Current`, `Previous` order
with colon-separated labels.
- Omit the `Previous` line for the first context window.
- Remove `<token_budget>` wrappers from newly rendered lineage,
near-compaction reminders, and `get_context_remaining` output.
- Keep recognizing legacy wrapped fragments so existing rollouts remain
compatible.
- Remove the post-sampling token snapshot that was only needed by the
periodic threshold path.

## Testing

- `just test -p codex-core token_budget` (11 tests passed)
This commit is contained in:
pakrym-oai
2026-06-20 21:50:09 -07:00
committed by GitHub
Unverified
parent 6df037d47f
commit aaf737fa59
7 changed files with 65 additions and 195 deletions
+17 -45
View File
@@ -3,63 +3,35 @@ use super::turn_context::TurnContext;
use crate::context::ContextualUserFragment;
use codex_features::Feature;
const TOKEN_BUDGET_USAGE_THRESHOLDS: [i64; 3] = [25, 50, 75];
pub(super) async fn maybe_record(
sess: &Session,
turn_context: &TurnContext,
tokens_before_sampling: i64,
tokens_after_sampling: i64,
tokens_until_compaction: i64,
) {
if !turn_context.config.features.enabled(Feature::TokenBudget) {
return;
}
let mut response_items = Vec::new();
if let Some(model_context_window) = turn_context.model_context_window()
&& model_context_window > 0
&& tokens_after_sampling > tokens_before_sampling
{
let tokens_before_sampling = tokens_before_sampling.max(0);
let tokens_after_sampling = tokens_after_sampling.max(0);
let crossed_threshold = TOKEN_BUDGET_USAGE_THRESHOLDS.iter().any(|threshold| {
tokens_before_sampling.saturating_mul(100)
< model_context_window.saturating_mul(*threshold)
&& tokens_after_sampling.saturating_mul(100)
>= model_context_window.saturating_mul(*threshold)
});
if crossed_threshold {
let tokens_left = model_context_window
.saturating_sub(tokens_after_sampling)
.max(0);
response_items.push(ContextualUserFragment::into(
crate::context::TokenBudgetRemainingContext::new(tokens_left),
));
}
}
if let Some(config) = turn_context.config.token_budget.as_ref().filter(|config| {
let Some(config) = turn_context.config.token_budget.as_ref().filter(|config| {
config
.reminder_threshold_tokens
.is_some_and(|threshold| tokens_until_compaction <= threshold)
}) {
let reminder_due = {
let mut state = sess.state.lock().await;
state.claim_token_budget_reminder()
};
if reminder_due {
response_items.push(ContextualUserFragment::into(
crate::context::TokenBudgetReminder::new(
&config.reminder_message_template,
tokens_until_compaction,
),
));
}
}) else {
return;
};
let reminder_due = {
let mut state = sess.state.lock().await;
state.claim_token_budget_reminder()
};
if !reminder_due {
return;
}
if !response_items.is_empty() {
sess.record_conversation_items(turn_context, &response_items)
.await;
}
let response_item = ContextualUserFragment::into(crate::context::TokenBudgetReminder::new(
&config.reminder_message_template,
tokens_until_compaction,
));
sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item))
.await;
}
+3 -12
View File
@@ -254,8 +254,7 @@ pub(crate) async fn run_turn(
window_id,
CodexResponsesRequestKind::Turn,
);
let tokens_before_sampling = sess.get_total_token_usage().await;
let (sampling_request_output, sampling_request_input) = run_sampling_request(
run_sampling_request(
Arc::clone(&sess),
Arc::clone(&turn_context),
Arc::clone(&turn_extension_data),
@@ -265,17 +264,11 @@ pub(crate) async fn run_turn(
sampling_request_input,
cancellation_token.child_token(),
)
.await?;
Ok((
tokens_before_sampling,
sampling_request_output,
sampling_request_input,
))
.await
}
.await;
match sampling_request_result {
Ok((tokens_before_sampling, sampling_request_output, sampling_request_input)) => {
Ok((sampling_request_output, sampling_request_input)) => {
let SamplingRequestResult {
needs_follow_up: model_needs_follow_up,
last_agent_message: sampling_request_last_agent_message,
@@ -326,8 +319,6 @@ pub(crate) async fn run_turn(
super::token_budget::maybe_record(
sess.as_ref(),
turn_context.as_ref(),
tokens_before_sampling,
tokens_after_sampling,
tokens_until_compaction,
)
.await;