From 6f33c9776370361206e3d582d69c2f50b0ed697f Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Sun, 13 Sep 2026 23:01:03 -0500 Subject: [PATCH] Fix provider context budgeting and ChatGPT subscription limits --- app/src/ai/blocklist/controller.rs | 98 +++++++++++++------ app/src/ai/blocklist/controller_tests.rs | 78 +++++++++++++++ app/src/ai/llms.rs | 24 ++--- app/src/ai/llms_tests.rs | 30 ++++-- app/src/ai/runtime/mod.rs | 3 +- .../ai/runtime/provider_run_coordinator.rs | 61 +++++++----- .../runtime/provider_run_coordinator_tests.rs | 57 ++++++++++- app/src/ai/runtime/rig.rs | 14 ++- app/src/ai/runtime/rig_tests.rs | 18 ++++ 9 files changed, 302 insertions(+), 81 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 25109a53..b1662372 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -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, ) -> 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(), diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 9637cc7b..b4269358 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -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); +} diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index d9798ea5..9bc98399 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -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, } } diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index 2a6c93a2..e182ff7b 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -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 + ); +} diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index b5e02dad..c3ea1b18 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -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; diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 91c430fa..75e9ac33 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -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) { 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( @@ -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::(); - 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::(); - 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> { diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index ced2b6f8..0a74d4d5 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -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); +} diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 438ddbcf..b5968e25 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -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, ) -> Option { - 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)] diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index bf4732ed..b23cfe4c 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -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) + ); + } +}