Fix provider context budgeting and ChatGPT subscription limits

This commit is contained in:
2026-09-13 23:01:03 -05:00
parent ba56d83fc6
commit 6f33c97763
9 changed files with 302 additions and 81 deletions
+67 -31
View File
@@ -81,11 +81,11 @@ use crate::ai::provider::ProviderConfig;
#[cfg(not(target_family = "wasm"))]
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
use crate::ai::runtime::{
prepare_provider_run, provider_runtime_for_request, OrchestrationModelOption,
PreparedProviderRun, ProviderActionContext, ProviderRunBlock, ProviderRunCoordinator,
ProviderRunProfile, ProviderRunProjection, ProviderRunResponseProjector,
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, RuntimeResponseConfig,
BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request,
OrchestrationModelOption, PreparedProviderRun, ProviderActionContext, ProviderRunBlock,
ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome,
RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
};
use crate::ai::AIRequestUsageModel;
use crate::cloud_object::model::persistence::CloudModel;
@@ -112,7 +112,7 @@ use crate::workspaces::user_workspaces::UserWorkspaces;
const PROGRESSIVE_SUMMARY_TRIGGER_USAGE: f32 = 0.70;
const PROGRESSIVE_SUMMARY_RETAINED_USAGE: f32 = 0.50;
const PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES: usize = 20;
const PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES: usize = 1;
const PROGRESSIVE_SUMMARY_OUTPUT_TOKENS: u32 = 8_000;
const PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS: u32 = 2_000;
const PROGRESSIVE_SUMMARY_START_TIMEOUT: Duration = Duration::from_secs(120);
@@ -199,7 +199,7 @@ fn progressive_summary_candidate(
}
fn estimated_text_tokens(text: &str) -> u32 {
u32::try_from(text.chars().count().div_ceil(4)).unwrap_or(u32::MAX)
u32::try_from(text.len().div_ceil(3)).unwrap_or(u32::MAX)
}
fn estimated_message_tokens(message: &ConversationMessage) -> u32 {
@@ -259,6 +259,8 @@ fn progressive_summary_split_point(
for index in (0..messages.len()).rev() {
let message_tokens = estimated_message_tokens(&messages[index]);
let retained_count = messages.len() - index;
// Keep the latest message, then retain only as much history as fits.
// A fixed message-count floor can retain many oversized file reads.
if retained_count > PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES
&& retained_tokens.saturating_add(message_tokens) > retained_token_budget
{
@@ -278,6 +280,17 @@ fn progressive_summary_split_point(
split_point -= 1;
}
// A single large tool batch can exceed the entire retained budget. Summarize
// the complete batch rather than keeping an oversized call/result pair forever.
// Its original results remain available in the tool-result archive.
if messages
.last()
.is_some_and(|message| content_contains_tool_result(&message.content))
&& estimated_history_tokens(&messages[split_point..]) > retained_token_budget
{
return messages.len();
}
split_point
}
@@ -5678,13 +5691,19 @@ impl BlocklistAIController {
self.fail_restored_provider_run(conversation_id, error, ctx);
return;
}
let base_provider_config =
ResponseStream::resolve_provider_config(snapshot.base_request.model.as_str(), ctx);
// A restored checkpoint may contain an older, larger catalog budget.
// Reconcile it before any resumed model call and persist the corrected limit.
snapshot.response_config.max_context_tokens = provider_context_window_tokens(
&base_provider_config,
snapshot.base_request.model.as_str(),
snapshot.response_config.max_context_tokens,
);
if let Err(error) = self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) {
self.fail_restored_provider_run(conversation_id, error, ctx);
return;
}
let base_provider_config =
ResponseStream::resolve_provider_config(snapshot.base_request.model.as_str(), ctx);
let cli_provider_config = snapshot
.cli_monitor_request
.as_ref()
@@ -6583,6 +6602,7 @@ impl BlocklistAIController {
fn begin_active_provider_progressive_summary(
&mut self,
conversation_id: AIConversationId,
force: bool,
ctx: &mut ModelContext<Self>,
) -> bool {
if self
@@ -6623,28 +6643,38 @@ impl BlocklistAIController {
return false;
};
(
conversation
.current_context_tokens()
.min(active_context_limit),
conversation.current_context_tokens(),
conversation.has_pending_progressive_summary(),
)
};
let transcript = active_run.coordinator.run().transcript();
let (estimated_context_tokens, persistent_context_tokens) = active_run
.coordinator
.estimated_context_tokens()
.unwrap_or_else(|| {
(
current_context_tokens,
current_context_tokens.saturating_sub(estimated_history_tokens(transcript)),
)
});
let current_context_tokens = current_context_tokens.max(estimated_context_tokens);
// Trigger and retain against the space history can actually use. Fixed
// prompts and a large output reservation cannot be reduced by summarizing.
let history_capacity = active_context_limit.saturating_sub(persistent_context_tokens);
let current_history_tokens =
current_context_tokens.saturating_sub(persistent_context_tokens);
if summary_pending
|| last_failure_context_tokens == Some(current_context_tokens)
|| current_context_tokens as f32 / (active_context_limit as f32)
< PROGRESSIVE_SUMMARY_TRIGGER_USAGE
|| (!force
&& current_history_tokens as f32 / (history_capacity.max(1) as f32)
< PROGRESSIVE_SUMMARY_TRIGGER_USAGE)
{
return false;
}
let transcript = active_run.coordinator.run().transcript();
let transcript_tokens = estimated_history_tokens(transcript);
let persistent_context_tokens = current_context_tokens.saturating_sub(transcript_tokens);
let retained_total_budget =
(active_context_limit as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32;
let retained_history_budget = retained_total_budget
.saturating_sub(persistent_context_tokens)
.saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS);
let retained_history_budget =
(history_capacity as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32;
let retained_history_budget =
retained_history_budget.saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_TOKENS);
let split_point = progressive_summary_split_point(transcript, retained_history_budget);
if split_point == 0 {
return false;
@@ -6851,6 +6881,16 @@ impl BlocklistAIController {
let mut summary_applied = false;
let mut recorded_usage = None;
match result {
Ok((summary, usage, summarizer_model_id))
if estimated_history_tokens(&progressive_summary_messages(summary.clone()))
>= estimated_history_tokens(&plan.summarized_prefix) =>
{
log::warn!(
"[progressive-summary] Summary did not reduce context for conversation {:?}",
conversation_id,
);
recorded_usage = Some((usage, summarizer_model_id));
}
Ok((summary, usage, summarizer_model_id)) => {
let Some(run) = self
.active_provider_runs
@@ -6925,10 +6965,6 @@ impl BlocklistAIController {
}
}
let failure_context_tokens = BlocklistAIHistoryModel::as_ref(ctx)
.conversation(&conversation_id)
.map(AIConversation::current_context_tokens)
.unwrap_or(plan.trigger_context_tokens);
if let Some(run) = self
.active_provider_runs
.get_mut(&conversation_id)
@@ -6938,7 +6974,7 @@ impl BlocklistAIController {
None
} else {
// Fail open and suppress another request at this unchanged context boundary.
Some(failure_context_tokens)
Some(plan.trigger_context_tokens)
};
}
if let Some((usage, summarizer_model_id)) = recorded_usage {
@@ -6998,7 +7034,7 @@ impl BlocklistAIController {
if self.apply_ready_active_provider_progressive_summary(conversation_id, ctx) {
return;
}
self.begin_active_provider_progressive_summary(conversation_id, ctx);
self.begin_active_provider_progressive_summary(conversation_id, false, ctx);
if matches!(
self.active_provider_progressive_summaries
.get(&conversation_id),
@@ -7471,7 +7507,7 @@ impl BlocklistAIController {
);
return;
}
if !self.begin_active_provider_progressive_summary(conversation_id, ctx) {
if !self.begin_active_provider_progressive_summary(conversation_id, true, ctx) {
self.fail_active_provider_run(
conversation_id,
"provider context overflow could not start compaction".to_string(),
+78
View File
@@ -3734,3 +3734,81 @@ fn optimistic_cli_subagent_completion_with_in_flight_stream_reports_success() {
});
});
}
#[test]
fn progressive_summary_compacts_short_histories_to_the_token_budget() {
let messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("large project context".repeat(1_000)),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::Text("Working on the requested change".to_string()),
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("Please continue".to_string()),
},
];
assert_eq!(super::progressive_summary_split_point(&messages, 100), 1);
}
#[test]
fn progressive_summary_includes_an_oversized_final_tool_batch() {
let messages = vec![
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "read".to_string(),
name: "read_file".to_string(),
input: serde_json::json!({"path": "large.rs"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "read".to_string(),
content: "file contents".repeat(1_000),
is_error: false,
},
},
];
assert_eq!(super::progressive_summary_split_point(&messages, 100), 2);
assert_eq!(
super::progressive_summary_split_point(&messages, 100_000),
0
);
}
#[test]
fn progressive_summary_keeps_retained_tool_results_paired() {
let messages = vec![
ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("old project context".repeat(1_000)),
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::ToolUse {
tool_use_id: "read".to_string(),
name: "read_file".to_string(),
input: serde_json::json!({"path": "small.rs"}),
},
},
ConversationMessage {
role: MessageRole::User,
content: MessageContent::ToolResult {
tool_use_id: "read".to_string(),
content: "file contents".to_string(),
is_error: false,
},
},
ConversationMessage {
role: MessageRole::Assistant,
content: MessageContent::Text("I found the issue".to_string()),
},
];
let budget = super::estimated_history_tokens(&messages[2..]);
assert_eq!(super::progressive_summary_split_point(&messages, budget), 1);
}
+13 -11
View File
@@ -2630,17 +2630,20 @@ fn get_new_agent_mode_choices(
#[cfg(not(target_family = "wasm"))]
fn openai_model_context_size(model: &OpenAIModelConfig, provider_kind: OpenAIProviderKind) -> u32 {
// Older saved ChatGPT catalogs stored the default window in max_input_tokens.
// Preserve their larger usable window until the catalog refresh separates them.
if provider_kind == OpenAIProviderKind::ChatGPTSubscription
&& model.default_context_size.is_none()
{
return model.context_size;
}
model
let input_limit = model
.max_input_tokens
.unwrap_or(model.context_size)
.min(model.context_size)
.min(model.context_size);
if provider_kind == OpenAIProviderKind::ChatGPTSubscription {
// The catalog's default window is the supported session budget. The
// advertised maximum must not silently opt subscriptions into long context.
// Older saved catalogs kept the default in max_input_tokens instead.
return model
.default_context_size
.unwrap_or(input_limit)
.min(input_limit);
}
input_limit
}
/// Merges endpoint metadata into a provider's configured models without
@@ -2749,8 +2752,7 @@ fn openai_model_context_window(
is_configurable: max_context_size > default_context_size,
min: default_context_size,
max: max_context_size,
// Use the full usable window; the catalog default remains available as
// the smaller configurable budget, not as a hard input ceiling.
// Subscription limits are already capped to the catalog's default above.
default_max: max_context_size,
}
}
+24 -6
View File
@@ -818,10 +818,10 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
&models[0],
crate::settings::OpenAIProviderKind::ChatGPTSubscription,
);
assert!(configurable.is_configurable);
assert!(!configurable.is_configurable);
assert_eq!(configurable.min, 258_400);
assert_eq!(configurable.default_max, 828_400);
assert_eq!(configurable.max, 828_400);
assert_eq!(configurable.default_max, 258_400);
assert_eq!(configurable.max, 258_400);
let fixed =
openai_model_context_window(&models[0], crate::settings::OpenAIProviderKind::LiteLLM);
@@ -883,7 +883,7 @@ fn provider_context_budget_respects_hard_input_and_total_limits() {
model.max_input_tokens = Some(400_000);
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
assert_eq!(window.min, 272_000);
assert_eq!(window.max, 400_000);
assert_eq!(window.max, 272_000);
assert_eq!(
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
window.max
@@ -891,13 +891,15 @@ fn provider_context_budget_respects_hard_input_and_total_limits() {
}
#[test]
fn legacy_chatgpt_default_does_not_limit_routing() {
fn legacy_chatgpt_default_limits_routing() {
let mut model = openai_model("legacy-model");
model.context_size = 872_000;
model.max_input_tokens = Some(272_000);
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
assert_eq!(window.min, 272_000);
assert_eq!(window.max, 872_000);
assert_eq!(window.max, 272_000);
assert_eq!(window.default_max, 272_000);
assert!(!window.is_configurable);
assert_eq!(
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
window.max
@@ -944,3 +946,19 @@ fn provider_registry_overrides_legacy_endpoint_for_duplicate_models() {
assert_eq!(config.api_key.as_deref(), Some("current-key"));
});
}
#[test]
fn chatgpt_default_window_limits_routing_without_a_separate_input_limit() {
let mut model = openai_model("subscription-model");
model.context_size = 872_000;
model.default_context_size = Some(272_000);
model.max_input_tokens = None;
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
assert_eq!(window.default_max, 272_000);
assert_eq!(window.max, 272_000);
assert!(!window.is_configurable);
assert_eq!(
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
272_000
);
}
+2 -1
View File
@@ -13,6 +13,7 @@ pub(crate) use provider_run_coordinator::{
CLI_MONITOR_PROVIDER_PROFILE,
};
pub(crate) use rig::{
prepare_provider_run, provider_runtime_for_request, PreparedProviderRun, ProviderActionContext,
prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request,
PreparedProviderRun, ProviderActionContext,
};
pub(crate) use rig_request::OrchestrationModelOption;
+35 -26
View File
@@ -244,6 +244,21 @@ impl ProviderRunCoordinator {
self.profiles.get(profile).map(|profile| &profile.request)
}
/// Includes the current transcript, fixed request context, tools, and output reservation.
/// Read at the model boundary so newly committed tool output is included.
pub(crate) fn estimated_context_tokens(&self) -> Option<(u32, u32)> {
let profile = self.profiles.get(self.run.profile().as_str())?;
let mut request = profile.request.clone();
request.prompt = None;
request.messages.clear();
let persistent_tokens = estimate_turn_request_tokens(&request);
request.messages = self.run.transcript().to_vec();
Some((
u32::try_from(estimate_turn_request_tokens(&request)).unwrap_or(u32::MAX),
u32::try_from(persistent_tokens).unwrap_or(u32::MAX),
))
}
pub(crate) fn set_max_context_tokens(&mut self, max_context_tokens: Option<u32>) {
self.max_context_tokens = max_context_tokens;
}
@@ -453,7 +468,10 @@ impl ProviderRunCoordinator {
return false;
};
let request = request_for_model_call(profile.request.clone(), call);
estimate_turn_request_tokens(&request) >= u64::from(max_context_tokens)
// Provider tokenizers and request envelopes differ. Compact before the
// estimate reaches the advertised limit, leaving room for that uncertainty.
let safe_context_tokens = u64::from(max_context_tokens) * 9 / 10;
estimate_turn_request_tokens(&request) >= safe_context_tokens
}
async fn checkpoint_or_fail<C>(
@@ -1120,42 +1138,33 @@ fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) -
template
}
const ESTIMATED_CHARS_PER_TOKEN: u64 = 4;
const ESTIMATED_BYTES_PER_TOKEN: u64 = 3;
const ESTIMATED_REQUEST_OVERHEAD_TOKENS: u64 = 512;
/// Deliberately overestimates request size without requiring provider-specific tokenizers.
/// Conservative estimate, not an exact tokenizer count. UTF-8 bytes account for
/// non-ASCII content better than character counts; the caller also leaves headroom.
fn estimate_turn_request_tokens(request: &TurnRequest) -> u64 {
let mut chars = request
let mut bytes = request
.system_prompt
.as_deref()
.map_or(0, |text| text.chars().count());
chars += request.prompt.as_ref().map_or(0, |prompt| {
serde_json::to_string(prompt)
.unwrap_or_default()
.chars()
.count()
.map_or(0, |text| text.len());
bytes += request.prompt.as_ref().map_or(0, |prompt| {
serde_json::to_string(prompt).unwrap_or_default().len()
});
chars += request
bytes += request
.messages
.iter()
.map(|message| {
serde_json::to_string(message)
.unwrap_or_default()
.chars()
.count()
})
.map(|message| serde_json::to_string(message).unwrap_or_default().len())
.sum::<usize>();
chars += request
bytes += request
.tools
.iter()
.map(|tool| {
serde_json::to_string(tool)
.unwrap_or_default()
.chars()
.count()
})
.map(|tool| serde_json::to_string(tool).unwrap_or_default().len())
.sum::<usize>();
let input_tokens = (chars as u64).div_ceil(ESTIMATED_CHARS_PER_TOKEN);
input_tokens.saturating_add(request.max_output_tokens.unwrap_or_default())
let input_tokens = (bytes as u64).div_ceil(ESTIMATED_BYTES_PER_TOKEN);
input_tokens
.saturating_add(ESTIMATED_REQUEST_OVERHEAD_TOKENS)
.saturating_add(request.max_output_tokens.unwrap_or_default())
}
fn tool_event_call_id(event: &ToolEvent) -> Result<&str, ProviderRunCoordinatorError> {
@@ -186,7 +186,7 @@ fn estimates_turn_request_with_system_tools_messages_and_output() {
+ serde_json::to_string(&request.tools[0]).unwrap().len()) as u64;
assert_eq!(
estimate_turn_request_tokens(&request),
expected_input.div_ceil(ESTIMATED_CHARS_PER_TOKEN) + 10
expected_input.div_ceil(ESTIMATED_BYTES_PER_TOKEN) + ESTIMATED_REQUEST_OVERHEAD_TOKENS + 10
);
}
@@ -1370,3 +1370,58 @@ fn session_runtime_is_rejected_before_any_turn_can_start() {
ProviderRunCoordinatorError::InvalidRuntime(_)
));
}
#[test]
fn context_estimate_tracks_live_transcript_and_reserves_fixed_context() {
let runtime = Arc::new(ScriptedRuntime::new(vec![]));
let mut request = request();
request.system_prompt = Some("project instructions".repeat(100));
request.max_output_tokens = Some(4_000);
let mut coordinator = ProviderRunCoordinator::from_request(
"context-budget",
runtime,
request,
Vec::new(),
ProviderRunLimits::default(),
)
.unwrap();
let (before, fixed_before) = coordinator.estimated_context_tokens().unwrap();
assert!(fixed_before > 4_000);
coordinator
.run_mut()
.compact_transcript_at_model_boundary(
1,
vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("new large tool output".repeat(1_000)),
}],
)
.unwrap();
let (after, fixed_after) = coordinator.estimated_context_tokens().unwrap();
assert!(after > before + 4_000);
assert_eq!(fixed_before, fixed_after);
}
#[tokio::test]
async fn preflight_leaves_headroom_before_the_advertised_context_limit() {
let runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()]));
let mut coordinator = coordinator(runtime.clone());
let (estimated_tokens, _) = coordinator.estimated_context_tokens().unwrap();
coordinator.set_max_context_tokens(Some(estimated_tokens + 1));
let (_sender, control) = turn_control();
let block = coordinator
.drive_until_blocked(control, |_| Ok(()))
.await
.unwrap();
assert_eq!(block, ProviderRunBlock::ContextWindowExceeded);
assert!(runtime.requests().is_empty());
}
#[test]
fn request_estimate_accounts_for_multibyte_text() {
let mut request = request();
request.system_prompt = Some("a".repeat(1_000));
let ascii_tokens = estimate_turn_request_tokens(&request);
request.system_prompt = Some("".repeat(1_000));
assert!(estimate_turn_request_tokens(&request) > ascii_tokens + 600);
}
+9 -5
View File
@@ -293,18 +293,22 @@ fn bedrock_max_output_tokens(model: &str) -> u64 {
.unwrap_or(64_000)
}
fn provider_context_window_tokens(
pub(crate) fn provider_context_window_tokens(
provider_config: &crate::ai::provider::ProviderConfig,
model: &str,
configured_limit: Option<u32>,
) -> Option<u32> {
configured_limit.or_else(|| match provider_config {
let provider_limit = match provider_config {
crate::ai::provider::ProviderConfig::Bedrock(_) => {
model_metadata(model).and_then(|metadata| metadata.context_window_tokens)
}
crate::ai::provider::ProviderConfig::OpenAI(_)
| crate::ai::provider::ProviderConfig::None => None,
})
crate::ai::provider::ProviderConfig::OpenAI(config) => config.max_input_tokens,
crate::ai::provider::ProviderConfig::None => None,
};
match (configured_limit, provider_limit) {
(Some(configured), Some(provider)) => Some(configured.min(provider)),
(configured, provider) => configured.or(provider),
}
}
#[cfg(test)]
+18
View File
@@ -55,3 +55,21 @@ async fn missing_cli_provider_route_falls_back_to_base_provider_profile() {
.iter()
.any(|tool| tool.name == "read_shell_command_output"));
}
#[test]
fn subscription_context_budget_uses_catalog_limit_and_clamps_stale_overrides() {
let mut config = openai_config("subscription-model");
config.kind = OpenAIProviderKind::ChatGPTSubscription;
config.max_input_tokens = Some(272_000);
let provider = ProviderConfig::OpenAI(config);
for (selected, expected) in [
(None, 272_000),
(Some(872_000), 272_000),
(Some(128_000), 128_000),
] {
assert_eq!(
super::provider_context_window_tokens(&provider, "subscription-model", selected),
Some(expected)
);
}
}