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"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
|
use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord};
|
||||||
use crate::ai::runtime::{
|
use crate::ai::runtime::{
|
||||||
prepare_provider_run, provider_runtime_for_request, OrchestrationModelOption,
|
prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request,
|
||||||
PreparedProviderRun, ProviderActionContext, ProviderRunBlock, ProviderRunCoordinator,
|
OrchestrationModelOption, PreparedProviderRun, ProviderActionContext, ProviderRunBlock,
|
||||||
ProviderRunProfile, ProviderRunProjection, ProviderRunResponseProjector,
|
ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection,
|
||||||
ProviderToolExecutionRef, ProviderToolLifecycleOutcome, RuntimeResponseConfig,
|
ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome,
|
||||||
BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
|
RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE,
|
||||||
};
|
};
|
||||||
use crate::ai::AIRequestUsageModel;
|
use crate::ai::AIRequestUsageModel;
|
||||||
use crate::cloud_object::model::persistence::CloudModel;
|
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_TRIGGER_USAGE: f32 = 0.70;
|
||||||
const PROGRESSIVE_SUMMARY_RETAINED_USAGE: f32 = 0.50;
|
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_OUTPUT_TOKENS: u32 = 8_000;
|
||||||
const PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS: u32 = 2_000;
|
const PROGRESSIVE_SUMMARY_CONTEXT_SAFETY_TOKENS: u32 = 2_000;
|
||||||
const PROGRESSIVE_SUMMARY_START_TIMEOUT: Duration = Duration::from_secs(120);
|
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 {
|
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 {
|
fn estimated_message_tokens(message: &ConversationMessage) -> u32 {
|
||||||
@@ -259,6 +259,8 @@ fn progressive_summary_split_point(
|
|||||||
for index in (0..messages.len()).rev() {
|
for index in (0..messages.len()).rev() {
|
||||||
let message_tokens = estimated_message_tokens(&messages[index]);
|
let message_tokens = estimated_message_tokens(&messages[index]);
|
||||||
let retained_count = messages.len() - 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
|
if retained_count > PROGRESSIVE_SUMMARY_MIN_RECENT_MESSAGES
|
||||||
&& retained_tokens.saturating_add(message_tokens) > retained_token_budget
|
&& retained_tokens.saturating_add(message_tokens) > retained_token_budget
|
||||||
{
|
{
|
||||||
@@ -278,6 +280,17 @@ fn progressive_summary_split_point(
|
|||||||
split_point -= 1;
|
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
|
split_point
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -5678,13 +5691,19 @@ impl BlocklistAIController {
|
|||||||
self.fail_restored_provider_run(conversation_id, error, ctx);
|
self.fail_restored_provider_run(conversation_id, error, ctx);
|
||||||
return;
|
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) {
|
if let Err(error) = self.persist_provider_run_snapshot(conversation_id, &snapshot, ctx) {
|
||||||
self.fail_restored_provider_run(conversation_id, error, ctx);
|
self.fail_restored_provider_run(conversation_id, error, ctx);
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
let base_provider_config =
|
|
||||||
ResponseStream::resolve_provider_config(snapshot.base_request.model.as_str(), ctx);
|
|
||||||
let cli_provider_config = snapshot
|
let cli_provider_config = snapshot
|
||||||
.cli_monitor_request
|
.cli_monitor_request
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -6583,6 +6602,7 @@ impl BlocklistAIController {
|
|||||||
fn begin_active_provider_progressive_summary(
|
fn begin_active_provider_progressive_summary(
|
||||||
&mut self,
|
&mut self,
|
||||||
conversation_id: AIConversationId,
|
conversation_id: AIConversationId,
|
||||||
|
force: bool,
|
||||||
ctx: &mut ModelContext<Self>,
|
ctx: &mut ModelContext<Self>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
if self
|
if self
|
||||||
@@ -6623,28 +6643,38 @@ impl BlocklistAIController {
|
|||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
(
|
(
|
||||||
conversation
|
conversation.current_context_tokens(),
|
||||||
.current_context_tokens()
|
|
||||||
.min(active_context_limit),
|
|
||||||
conversation.has_pending_progressive_summary(),
|
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
|
if summary_pending
|
||||||
|| last_failure_context_tokens == Some(current_context_tokens)
|
|| last_failure_context_tokens == Some(current_context_tokens)
|
||||||
|| current_context_tokens as f32 / (active_context_limit as f32)
|
|| (!force
|
||||||
< PROGRESSIVE_SUMMARY_TRIGGER_USAGE
|
&& current_history_tokens as f32 / (history_capacity.max(1) as f32)
|
||||||
|
< PROGRESSIVE_SUMMARY_TRIGGER_USAGE)
|
||||||
{
|
{
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
let retained_history_budget =
|
||||||
let transcript = active_run.coordinator.run().transcript();
|
(history_capacity as f32 * PROGRESSIVE_SUMMARY_RETAINED_USAGE) as u32;
|
||||||
let transcript_tokens = estimated_history_tokens(transcript);
|
let retained_history_budget =
|
||||||
let persistent_context_tokens = current_context_tokens.saturating_sub(transcript_tokens);
|
retained_history_budget.saturating_sub(PROGRESSIVE_SUMMARY_OUTPUT_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 split_point = progressive_summary_split_point(transcript, retained_history_budget);
|
let split_point = progressive_summary_split_point(transcript, retained_history_budget);
|
||||||
if split_point == 0 {
|
if split_point == 0 {
|
||||||
return false;
|
return false;
|
||||||
@@ -6851,6 +6881,16 @@ impl BlocklistAIController {
|
|||||||
let mut summary_applied = false;
|
let mut summary_applied = false;
|
||||||
let mut recorded_usage = None;
|
let mut recorded_usage = None;
|
||||||
match result {
|
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)) => {
|
Ok((summary, usage, summarizer_model_id)) => {
|
||||||
let Some(run) = self
|
let Some(run) = self
|
||||||
.active_provider_runs
|
.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
|
if let Some(run) = self
|
||||||
.active_provider_runs
|
.active_provider_runs
|
||||||
.get_mut(&conversation_id)
|
.get_mut(&conversation_id)
|
||||||
@@ -6938,7 +6974,7 @@ impl BlocklistAIController {
|
|||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
// Fail open and suppress another request at this unchanged context boundary.
|
// 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 {
|
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) {
|
if self.apply_ready_active_provider_progressive_summary(conversation_id, ctx) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
self.begin_active_provider_progressive_summary(conversation_id, ctx);
|
self.begin_active_provider_progressive_summary(conversation_id, false, ctx);
|
||||||
if matches!(
|
if matches!(
|
||||||
self.active_provider_progressive_summaries
|
self.active_provider_progressive_summaries
|
||||||
.get(&conversation_id),
|
.get(&conversation_id),
|
||||||
@@ -7471,7 +7507,7 @@ impl BlocklistAIController {
|
|||||||
);
|
);
|
||||||
return;
|
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(
|
self.fail_active_provider_run(
|
||||||
conversation_id,
|
conversation_id,
|
||||||
"provider context overflow could not start compaction".to_string(),
|
"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"))]
|
#[cfg(not(target_family = "wasm"))]
|
||||||
fn openai_model_context_size(model: &OpenAIModelConfig, provider_kind: OpenAIProviderKind) -> u32 {
|
fn openai_model_context_size(model: &OpenAIModelConfig, provider_kind: OpenAIProviderKind) -> u32 {
|
||||||
// Older saved ChatGPT catalogs stored the default window in max_input_tokens.
|
let input_limit = model
|
||||||
// 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
|
|
||||||
.max_input_tokens
|
.max_input_tokens
|
||||||
.unwrap_or(model.context_size)
|
.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
|
/// 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,
|
is_configurable: max_context_size > default_context_size,
|
||||||
min: default_context_size,
|
min: default_context_size,
|
||||||
max: max_context_size,
|
max: max_context_size,
|
||||||
// Use the full usable window; the catalog default remains available as
|
// Subscription limits are already capped to the catalog's default above.
|
||||||
// the smaller configurable budget, not as a hard input ceiling.
|
|
||||||
default_max: max_context_size,
|
default_max: max_context_size,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -818,10 +818,10 @@ fn chatgpt_codex_models_parse_visible_catalog_entries() {
|
|||||||
&models[0],
|
&models[0],
|
||||||
crate::settings::OpenAIProviderKind::ChatGPTSubscription,
|
crate::settings::OpenAIProviderKind::ChatGPTSubscription,
|
||||||
);
|
);
|
||||||
assert!(configurable.is_configurable);
|
assert!(!configurable.is_configurable);
|
||||||
assert_eq!(configurable.min, 258_400);
|
assert_eq!(configurable.min, 258_400);
|
||||||
assert_eq!(configurable.default_max, 828_400);
|
assert_eq!(configurable.default_max, 258_400);
|
||||||
assert_eq!(configurable.max, 828_400);
|
assert_eq!(configurable.max, 258_400);
|
||||||
|
|
||||||
let fixed =
|
let fixed =
|
||||||
openai_model_context_window(&models[0], crate::settings::OpenAIProviderKind::LiteLLM);
|
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);
|
model.max_input_tokens = Some(400_000);
|
||||||
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
|
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
|
||||||
assert_eq!(window.min, 272_000);
|
assert_eq!(window.min, 272_000);
|
||||||
assert_eq!(window.max, 400_000);
|
assert_eq!(window.max, 272_000);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
|
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
|
||||||
window.max
|
window.max
|
||||||
@@ -891,13 +891,15 @@ fn provider_context_budget_respects_hard_input_and_total_limits() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn legacy_chatgpt_default_does_not_limit_routing() {
|
fn legacy_chatgpt_default_limits_routing() {
|
||||||
let mut model = openai_model("legacy-model");
|
let mut model = openai_model("legacy-model");
|
||||||
model.context_size = 872_000;
|
model.context_size = 872_000;
|
||||||
model.max_input_tokens = Some(272_000);
|
model.max_input_tokens = Some(272_000);
|
||||||
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
|
let window = openai_model_context_window(&model, OpenAIProviderKind::ChatGPTSubscription);
|
||||||
assert_eq!(window.min, 272_000);
|
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!(
|
assert_eq!(
|
||||||
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
|
openai_model_context_size(&model, OpenAIProviderKind::ChatGPTSubscription),
|
||||||
window.max
|
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"));
|
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,
|
CLI_MONITOR_PROVIDER_PROFILE,
|
||||||
};
|
};
|
||||||
pub(crate) use rig::{
|
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;
|
pub(crate) use rig_request::OrchestrationModelOption;
|
||||||
|
|||||||
@@ -244,6 +244,21 @@ impl ProviderRunCoordinator {
|
|||||||
self.profiles.get(profile).map(|profile| &profile.request)
|
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>) {
|
pub(crate) fn set_max_context_tokens(&mut self, max_context_tokens: Option<u32>) {
|
||||||
self.max_context_tokens = max_context_tokens;
|
self.max_context_tokens = max_context_tokens;
|
||||||
}
|
}
|
||||||
@@ -453,7 +468,10 @@ impl ProviderRunCoordinator {
|
|||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
let request = request_for_model_call(profile.request.clone(), call);
|
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>(
|
async fn checkpoint_or_fail<C>(
|
||||||
@@ -1120,42 +1138,33 @@ fn request_for_model_call(mut template: TurnRequest, call: &ProviderModelCall) -
|
|||||||
template
|
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 {
|
fn estimate_turn_request_tokens(request: &TurnRequest) -> u64 {
|
||||||
let mut chars = request
|
let mut bytes = request
|
||||||
.system_prompt
|
.system_prompt
|
||||||
.as_deref()
|
.as_deref()
|
||||||
.map_or(0, |text| text.chars().count());
|
.map_or(0, |text| text.len());
|
||||||
chars += request.prompt.as_ref().map_or(0, |prompt| {
|
bytes += request.prompt.as_ref().map_or(0, |prompt| {
|
||||||
serde_json::to_string(prompt)
|
serde_json::to_string(prompt).unwrap_or_default().len()
|
||||||
.unwrap_or_default()
|
|
||||||
.chars()
|
|
||||||
.count()
|
|
||||||
});
|
});
|
||||||
chars += request
|
bytes += request
|
||||||
.messages
|
.messages
|
||||||
.iter()
|
.iter()
|
||||||
.map(|message| {
|
.map(|message| serde_json::to_string(message).unwrap_or_default().len())
|
||||||
serde_json::to_string(message)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.chars()
|
|
||||||
.count()
|
|
||||||
})
|
|
||||||
.sum::<usize>();
|
.sum::<usize>();
|
||||||
chars += request
|
bytes += request
|
||||||
.tools
|
.tools
|
||||||
.iter()
|
.iter()
|
||||||
.map(|tool| {
|
.map(|tool| serde_json::to_string(tool).unwrap_or_default().len())
|
||||||
serde_json::to_string(tool)
|
|
||||||
.unwrap_or_default()
|
|
||||||
.chars()
|
|
||||||
.count()
|
|
||||||
})
|
|
||||||
.sum::<usize>();
|
.sum::<usize>();
|
||||||
let input_tokens = (chars as u64).div_ceil(ESTIMATED_CHARS_PER_TOKEN);
|
let input_tokens = (bytes as u64).div_ceil(ESTIMATED_BYTES_PER_TOKEN);
|
||||||
input_tokens.saturating_add(request.max_output_tokens.unwrap_or_default())
|
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> {
|
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;
|
+ serde_json::to_string(&request.tools[0]).unwrap().len()) as u64;
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
estimate_turn_request_tokens(&request),
|
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(_)
|
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)
|
.unwrap_or(64_000)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn provider_context_window_tokens(
|
pub(crate) fn provider_context_window_tokens(
|
||||||
provider_config: &crate::ai::provider::ProviderConfig,
|
provider_config: &crate::ai::provider::ProviderConfig,
|
||||||
model: &str,
|
model: &str,
|
||||||
configured_limit: Option<u32>,
|
configured_limit: Option<u32>,
|
||||||
) -> Option<u32> {
|
) -> Option<u32> {
|
||||||
configured_limit.or_else(|| match provider_config {
|
let provider_limit = match provider_config {
|
||||||
crate::ai::provider::ProviderConfig::Bedrock(_) => {
|
crate::ai::provider::ProviderConfig::Bedrock(_) => {
|
||||||
model_metadata(model).and_then(|metadata| metadata.context_window_tokens)
|
model_metadata(model).and_then(|metadata| metadata.context_window_tokens)
|
||||||
}
|
}
|
||||||
crate::ai::provider::ProviderConfig::OpenAI(_)
|
crate::ai::provider::ProviderConfig::OpenAI(config) => config.max_input_tokens,
|
||||||
| crate::ai::provider::ProviderConfig::None => None,
|
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)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -55,3 +55,21 @@ async fn missing_cli_provider_route_falls_back_to_base_provider_profile() {
|
|||||||
.iter()
|
.iter()
|
||||||
.any(|tool| tool.name == "read_shell_command_output"));
|
.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