Support model-owned token budget defaults (#35608)

## What changed

- Add token-budget settings to model catalog messages and apply them when the
  feature is enabled without explicit token-budget configuration.
- Keep explicit user settings authoritative, reject invalid catalog defaults,
  and preserve resolved defaults in exported config locks.
- Manage context-window guidance through world state so it updates once when
  the active model changes while retaining prior conversation history.

## Testing

- Cover catalog defaults, explicit overrides, disabled features, invalid
  values, config-lock replay, and model switching.

GitOrigin-RevId: 54544fefaa14b09f5ef5bad9967a13c52b87b0fe
This commit is contained in:
rka-oai
2026-07-27 12:52:37 +00:00
committed by copyberry
parent 16d7daad7c
commit fbe65995bb
17 changed files with 742 additions and 72 deletions
+76 -60
View File
@@ -1120,6 +1120,78 @@ pub struct TokenBudgetConfig {
}
impl TokenBudgetConfig {
pub(crate) fn validate(&self) -> std::io::Result<()> {
if self
.reminder_threshold_tokens
.is_some_and(|tokens| tokens <= 0)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.reminder_threshold_tokens must be positive",
));
}
if self.reminder_message_template.trim().is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.reminder_message_template must not be empty",
));
}
if self.reminder_message_template.len() > TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"features.token_budget.reminder_message_template must not exceed {TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES} bytes"
),
));
}
if self
.guidance_message
.as_ref()
.is_some_and(|message| message.len() > TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"features.token_budget.guidance_message must not exceed {TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES} bytes"
),
));
}
if self
.auto_compact_fallback_prompt
.as_ref()
.is_some_and(|prompt| prompt.len() > AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"features.token_budget.auto_compact_fallback_prompt must not exceed {AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES} bytes"
),
));
}
if self.auto_compact_fallback_prompt.is_some()
&& self.auto_compact_fallback_buffer_tokens.is_none()
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.auto_compact_fallback_buffer_tokens is required when auto_compact_fallback_prompt is set",
));
}
if self
.auto_compact_fallback_buffer_tokens
.is_some_and(|tokens| tokens <= 0)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.auto_compact_fallback_buffer_tokens must be positive",
));
}
Ok(())
}
pub(crate) fn fallback_buffer_tokens(&self) -> i64 {
if self.auto_compact_fallback_prompt.is_some() {
self.auto_compact_fallback_buffer_tokens.unwrap_or(0)
@@ -2684,85 +2756,29 @@ fn resolve_token_budget_config(
let token_budget_config = token_budget_toml_config(config_toml.features.as_ref());
let reminder_threshold_tokens =
token_budget_config.and_then(|config| config.reminder_threshold_tokens);
if reminder_threshold_tokens.is_some_and(|tokens| tokens <= 0) {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.reminder_threshold_tokens must be positive",
));
}
let reminder_message_template = token_budget_config
.and_then(|config| config.reminder_message_template.clone())
.unwrap_or_else(|| DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string());
if reminder_message_template.trim().is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.reminder_message_template must not be empty",
));
}
if reminder_message_template.len() > TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"features.token_budget.reminder_message_template must not exceed {TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES} bytes"
),
));
}
let guidance_message = token_budget_config
.and_then(|config| config.guidance_message.clone())
.filter(|message| !message.trim().is_empty());
if guidance_message
.as_ref()
.is_some_and(|message| message.len() > TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"features.token_budget.guidance_message must not exceed {TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES} bytes"
),
));
}
let auto_compact_fallback_prompt = token_budget_config
.and_then(|config| config.auto_compact_fallback_prompt.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string);
if auto_compact_fallback_prompt
.as_ref()
.is_some_and(|prompt| prompt.len() > AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES)
{
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"features.token_budget.auto_compact_fallback_prompt must not exceed {AUTO_COMPACT_FALLBACK_PROMPT_MAX_BYTES} bytes"
),
));
}
let auto_compact_fallback_buffer_tokens =
token_budget_config.and_then(|config| config.auto_compact_fallback_buffer_tokens);
if auto_compact_fallback_prompt.is_some() && auto_compact_fallback_buffer_tokens.is_none() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.auto_compact_fallback_buffer_tokens is required when auto_compact_fallback_prompt is set",
));
}
if auto_compact_fallback_buffer_tokens.is_some_and(|tokens| tokens <= 0) {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"features.token_budget.auto_compact_fallback_buffer_tokens must be positive",
));
}
Ok(Some(TokenBudgetConfig {
let token_budget = TokenBudgetConfig {
reminder_threshold_tokens,
reminder_message_template,
guidance_message,
auto_compact_fallback_prompt,
auto_compact_fallback_buffer_tokens,
}))
};
token_budget.validate()?;
Ok(Some(token_budget))
}
fn resolve_rollout_budget_config(