public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix provider context budgeting and ChatGPT subscription limits

Browse Source
This commit is contained in:
rkw6086
2026-09-13 23:01:03 -05:00
parent ba56d83fc6
commit 6f33c97763
9 changed files with 302 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/runtime/mod.rs
+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;
app/src/ai/runtime/provider_run_coordinator.rs
+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> {
app/src/ai/runtime/provider_run_coordinator_tests.rs
+56 -1
View File
@@ -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);
}
app/src/ai/runtime/rig.rs
+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)]
app/src/ai/runtime/rig_tests.rs
+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)
);
}
}