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

[codex] add configurable token budget compaction reminder (#29255)

Browse Source 
## Why

The token-budget feature reports coarse remaining-context milestones,
but it does not give the model a configurable wrap-up prompt before
automatic compaction. A strict threshold-crossing check can also miss
resumed or reconfigured windows that are already inside the threshold.

## What changed

- Add structured `[features.token_budget]` configuration for an absolute
`reminder_threshold_tokens` and bounded `reminder_message_template`;
`{n_remaining}` is expanded when the reminder is delivered.
- Compute remaining tokens against the next effective auto-compaction
boundary, including scoped `body_after_prefix` accounting and the full
context-window limit.
- Make reminder delivery level-triggered before and after sampling, with
one-shot state owned by `AutoCompactWindow` and re-armed on compaction,
`new_context`, restore, or history replacement.
- Leave the existing initial full-window token-budget context, 25/50/75%
notices, and token-budget tools unchanged.
- Persist the resolved feature configuration in the session config lock
and regenerate the config schema.

## Validation

- `just test -p codex-core token_budget`
- `just test -p codex-core
token_budget_reminder_emits_after_crossing_compaction_threshold`
- `just test -p codex-core auto_compact_window`
- `just test -p codex-core
lock_contains_prompts_and_materializes_features`
- `just test -p codex-features`
- `just test -p codex-config`
This commit is contained in:
pakrym-oai
2026-06-20 19:13:42 -07:00
committed by GitHub
Unverified
parent b6d6be2a84
commit 6df037d47f
15 changed files with 418 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/config_lock.rs
+26
View File
@@ -10,6 +10,7 @@ use codex_features::FeatureToml;
use codex_features::FeaturesToml;
use codex_features::MultiAgentV2ConfigToml;
use codex_features::RolloutBudgetConfigToml;
use codex_features::TokenBudgetConfigToml;
use codex_protocol::ThreadId;
use crate::config::Config;
@@ -152,6 +153,12 @@ fn save_config_resolved_fields(
resolved_config_to_toml(&config.multi_agent_v2, "features.multi_agent_v2")?;
multi_agent_v2.enabled = Some(config.features.enabled(Feature::MultiAgentV2));
features.multi_agent_v2 = Some(FeatureToml::Config(multi_agent_v2));
if let Some(token_budget) = config.token_budget.as_ref() {
let mut token_budget: TokenBudgetConfigToml =
resolved_config_to_toml(token_budget, "features.token_budget")?;
token_budget.enabled = Some(config.features.enabled(Feature::TokenBudget));
features.token_budget = Some(FeatureToml::Config(token_budget));
}
if let Some(rollout_budget) = config.rollout_budget.as_ref() {
let mut rollout_budget: RolloutBudgetConfigToml =
resolved_config_to_toml(rollout_budget, "features.rollout_budget")?;
@@ -238,6 +245,14 @@ mod tests {
async fn lock_contains_prompts_and_materializes_features() {
let mut sc = crate::session::tests::make_session_configuration_for_tests().await;
let mut config = (*sc.original_config_do_not_use).clone();
config.token_budget = Some(crate::config::TokenBudgetConfig {
reminder_threshold_tokens: Some(16_000),
reminder_message_template: "Locked reminder: {n_remaining} tokens.".to_string(),
});
config
.features
.enable(Feature::TokenBudget)
.expect("token_budget should be enableable in tests");
config.rollout_budget = Some(crate::config::RolloutBudgetConfig {
limit_tokens: 100_000,
reminder_interval_tokens: 10_000,
@@ -318,6 +333,17 @@ mod tests {
})
));
assert_eq!(
features.token_budget,
Some(FeatureToml::Config(TokenBudgetConfigToml {
enabled: Some(true),
reminder_threshold_tokens: Some(16_000),
reminder_message_template: Some(
"Locked reminder: {n_remaining} tokens.".to_string()
),
}))
);
assert_eq!(
features.rollout_budget,
Some(FeatureToml::Config(RolloutBudgetConfigToml {
codex-rs/core/src/session/token_budget.rs
+45 -24
View File
@@ -5,40 +5,61 @@ use codex_features::Feature;
const TOKEN_BUDGET_USAGE_THRESHOLDS: [i64; 3] = [25, 50, 75];
pub(super) async fn maybe_record_token_budget_remaining_context(
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 Some(model_context_window) = turn_context.model_context_window() else {
return;
};
if model_context_window <= 0 || tokens_after_sampling <= tokens_before_sampling {
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),
));
}
}
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 {
return;
if 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,
),
));
}
}
let tokens_left = model_context_window
.saturating_sub(tokens_after_sampling)
.max(0);
let response_item = ContextualUserFragment::into(
crate::context::TokenBudgetRemainingContext::new(tokens_left),
);
sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item))
.await;
if !response_items.is_empty() {
sess.record_conversation_items(turn_context, &response_items)
.await;
}
}
codex-rs/core/src/session/turn.rs
+12 -1
View File
@@ -313,11 +313,22 @@ pub(crate) async fn run_turn(
);
let tokens_after_sampling = token_status.active_context_tokens;
super::token_budget::maybe_record_token_budget_remaining_context(
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);
super::token_budget::maybe_record(
sess.as_ref(),
turn_context.as_ref(),
tokens_before_sampling,
tokens_after_sampling,
tokens_until_compaction,
)
.await;