Fix provider context budgeting and ChatGPT subscription limits
This commit is contained in:
@@ -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(),
|
||||
|
||||
@@ -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
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user