Skip to content

Commit 92116a5

Browse files
committed
fix(ai): derive safe model output token limits
Hide output-token configuration from the model form while preserving existing explicit overrides. Derive a shared output limit for model requests and context compression when the value is absent or invalid. Reject context windows below 32K and fall back with a warning when an explicit output limit exceeds 40% of the context window. Refs #1558
1 parent 5a089bf commit 92116a5

8 files changed

Lines changed: 194 additions & 45 deletions

File tree

src/crates/assembly/core/src/agentic/execution/execution_engine.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ use crate::agentic::WorkspaceBinding;
3939
use crate::infrastructure::ai::get_global_ai_client_factory;
4040
use crate::service::config::get_global_config_service;
4141
use crate::service::config::types::{
42-
model_runtime_binding_fingerprint, ModelCapability, ModelCategory,
42+
automatic_max_output_tokens, model_runtime_binding_fingerprint, ModelCapability, ModelCategory,
4343
};
4444
use crate::util::errors::{BitFunError, BitFunResult};
4545
use crate::util::token_counter::TokenCounter;
@@ -342,7 +342,6 @@ pub struct ExecutionEngine {
342342
}
343343

344344
impl ExecutionEngine {
345-
const AUTO_COMPRESSION_DEFAULT_OUTPUT_RESERVE_TOKENS: usize = 16_000;
346345
const AUTO_COMPRESSION_SAFETY_RESERVE_TOKENS: usize = 10_000;
347346
const FINALIZE_AFTER_REPEATED_TOOL_FAILURES_REMINDER: &'static str = "This turn must end now because repeated tool failures have prevented further progress. Ignore any unfinished work. Your task now is to give the user a final answer. Do not call any more tools; any tool call will fail. Respond in plain text only. Summarize what was completed, what failed, the evidence available from the tool results, and the single best next step for the user.";
348347
const FINALIZE_AFTER_MAX_ROUNDS_REMINDER: &'static str = "This turn must end now because it has reached the round limit. Ignore any unfinished work. Your task now is to give the user a final answer. Do not call any more tools; any tool call will fail. Respond in plain text only. Summarize the most useful completed work and evidence collected so far, and clearly distinguish resolved items from anything still unresolved.";
@@ -516,7 +515,7 @@ impl ExecutionEngine {
516515
) -> CompressionTriggerBudget {
517516
let output_reserve_tokens = configured_max_tokens
518517
.map(|value| value as usize)
519-
.unwrap_or(Self::AUTO_COMPRESSION_DEFAULT_OUTPUT_RESERVE_TOKENS);
518+
.unwrap_or_else(|| automatic_max_output_tokens(context_window as u32) as usize);
520519
let safety_reserve_tokens = Self::AUTO_COMPRESSION_SAFETY_RESERVE_TOKENS;
521520
let input_limit =
522521
context_window.saturating_sub(output_reserve_tokens + safety_reserve_tokens);
@@ -4074,12 +4073,12 @@ mod tests {
40744073
}
40754074

40764075
#[test]
4077-
fn compression_trigger_budget_uses_16k_output_reserve_when_max_tokens_is_unset() {
4076+
fn compression_trigger_budget_uses_the_automatic_output_tier_when_max_tokens_is_unset() {
40784077
let budget = ExecutionEngine::compression_trigger_budget(128_000, None);
40794078

4080-
assert_eq!(budget.output_reserve_tokens, 16_000);
4079+
assert_eq!(budget.output_reserve_tokens, 32_000);
40814080
assert_eq!(budget.safety_reserve_tokens, 10_000);
4082-
assert_eq!(budget.input_limit, 102_000);
4081+
assert_eq!(budget.input_limit, 86_000);
40834082
}
40844083

40854084
#[test]

src/crates/assembly/core/src/service/config/providers.rs

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,10 +72,10 @@ impl ConfigProvider for AIConfigProvider {
7272
warnings.push(format!("Model '{}' has empty API key", model.name));
7373
}
7474
if let Some(context_window) = model.context_window {
75-
if context_window == 0 {
75+
if context_window < MIN_MODEL_CONTEXT_WINDOW_TOKENS {
7676
return Err(BitFunError::validation(format!(
77-
"Model '{}' context_window must be greater than 0",
78-
model.name
77+
"Model '{}' context_window must be at least {}",
78+
model.name, MIN_MODEL_CONTEXT_WINDOW_TOKENS
7979
)));
8080
}
8181
}
@@ -160,6 +160,32 @@ impl ConfigProvider for AIConfigProvider {
160160
}
161161
}
162162

163+
#[cfg(test)]
164+
mod tests {
165+
use super::*;
166+
167+
#[tokio::test]
168+
async fn rejects_a_model_context_window_smaller_than_32k() {
169+
let mut config = AIConfig::default();
170+
config.models.push(AIModelConfig {
171+
name: "Test model".to_string(),
172+
provider: "openai".to_string(),
173+
context_window: Some(MIN_MODEL_CONTEXT_WINDOW_TOKENS - 1),
174+
..AIModelConfig::default()
175+
});
176+
let value = serde_json::to_value(config).expect("AI config should serialize");
177+
178+
let error = AIConfigProvider
179+
.validate_config(&value)
180+
.await
181+
.expect_err("small context windows must be rejected");
182+
183+
assert!(error
184+
.to_string()
185+
.contains("context_window must be at least 32000"));
186+
}
187+
}
188+
163189
/// Theme system configuration provider (new, supports theme management).
164190
pub struct ThemesConfigProvider;
165191

src/crates/assembly/core/src/service/config/types.rs

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1222,6 +1222,30 @@ pub enum AgentSubagentOverrideState {
12221222
pub type ParentSubagentOverrideConfig = HashMap<String, AgentSubagentOverrideState>;
12231223
pub type AgentSubagentOverrideConfig = HashMap<String, ParentSubagentOverrideConfig>;
12241224

1225+
pub const DEFAULT_MODEL_CONTEXT_WINDOW_TOKENS: u32 = 128_128;
1226+
pub const MIN_MODEL_CONTEXT_WINDOW_TOKENS: u32 = 32_000;
1227+
pub const MAX_CONFIGURED_OUTPUT_TOKENS_RATIO_PERCENT: u32 = 40;
1228+
const AUTOMATIC_MAX_OUTPUT_TOKEN_TIERS: [u32; 5] = [8_000, 16_000, 24_000, 32_000, 64_000];
1229+
1230+
/// Chooses the largest supported output tier that does not exceed one quarter
1231+
/// of the model context window.
1232+
pub fn automatic_max_output_tokens(context_window: u32) -> u32 {
1233+
let quarter_context = context_window / 4;
1234+
AUTOMATIC_MAX_OUTPUT_TOKEN_TIERS
1235+
.iter()
1236+
.rev()
1237+
.copied()
1238+
.find(|tier| *tier <= quarter_context)
1239+
.unwrap_or(quarter_context)
1240+
}
1241+
1242+
/// A configured output cap may use up to 40% of the model context window.
1243+
pub fn is_valid_configured_max_output_tokens(context_window: u32, max_tokens: u32) -> bool {
1244+
max_tokens > 0
1245+
&& u64::from(max_tokens) * 100
1246+
<= u64::from(context_window) * u64::from(MAX_CONFIGURED_OUTPUT_TOKENS_RATIO_PERCENT)
1247+
}
1248+
12251249
#[derive(Debug, Clone, Serialize, Deserialize)]
12261250
#[serde(default, from = "AIModelConfigCompat")]
12271251
pub struct AIModelConfig {
@@ -1239,7 +1263,8 @@ pub struct AIModelConfig {
12391263
pub api_key: String,
12401264
/// Context window size (total token limit for input + output).
12411265
pub context_window: Option<u32>,
1242-
/// Max output tokens (request parameter limiting model output length).
1266+
/// Optional advanced override for the request output limit. When absent,
1267+
/// BitFun derives a tiered limit from the context window at runtime.
12431268
pub max_tokens: Option<u32>,
12441269
pub temperature: Option<f64>,
12451270
pub top_p: Option<f64>,

src/crates/assembly/core/src/util/types/config.rs

Lines changed: 95 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
11
use crate::service::config::types::AIModelConfig;
2+
use crate::service::config::types::{
3+
automatic_max_output_tokens, is_valid_configured_max_output_tokens,
4+
DEFAULT_MODEL_CONTEXT_WINDOW_TOKENS, MAX_CONFIGURED_OUTPUT_TOKENS_RATIO_PERCENT,
5+
MIN_MODEL_CONTEXT_WINDOW_TOKENS,
6+
};
27
pub use bitfun_core_types::AIConfig;
38
use log::warn;
49

@@ -94,15 +99,50 @@ impl TryFrom<AIModelConfig> for AIConfig {
9499
resolve_request_url(&other.base_url, &other.provider, &other.model_name)
95100
});
96101

102+
let context_window = other
103+
.context_window
104+
.unwrap_or(DEFAULT_MODEL_CONTEXT_WINDOW_TOKENS);
105+
if context_window < MIN_MODEL_CONTEXT_WINDOW_TOKENS {
106+
return Err(format!(
107+
"Model '{}' context_window must be at least {}",
108+
other.name, MIN_MODEL_CONTEXT_WINDOW_TOKENS
109+
));
110+
}
111+
112+
let max_tokens = match other.max_tokens {
113+
Some(configured_max_tokens)
114+
if is_valid_configured_max_output_tokens(context_window, configured_max_tokens) =>
115+
{
116+
configured_max_tokens
117+
}
118+
Some(configured_max_tokens) => {
119+
let automatic_max_tokens = automatic_max_output_tokens(context_window);
120+
let maximum_allowed_tokens = u64::from(context_window)
121+
* u64::from(MAX_CONFIGURED_OUTPUT_TOKENS_RATIO_PERCENT)
122+
/ 100;
123+
warn!(
124+
"Invalid model max_tokens; falling back to automatic output limit: model_id={}, model_name={}, context_window={}, configured_max_tokens={}, maximum_allowed_tokens={}, automatic_max_tokens={}",
125+
other.id,
126+
other.name,
127+
context_window,
128+
configured_max_tokens,
129+
maximum_allowed_tokens,
130+
automatic_max_tokens
131+
);
132+
automatic_max_tokens
133+
}
134+
None => automatic_max_output_tokens(context_window),
135+
};
136+
97137
Ok(AIConfig {
98138
name: other.name.clone(),
99139
base_url: other.base_url.clone(),
100140
request_url,
101141
api_key: other.api_key.clone(),
102142
model: other.model_name.clone(),
103143
format: other.provider.clone(),
104-
context_window: other.context_window.unwrap_or(128128),
105-
max_tokens: other.max_tokens,
144+
context_window,
145+
max_tokens: Some(max_tokens),
106146
temperature: other.temperature,
107147
top_p: other.top_p,
108148
reasoning_mode,
@@ -237,4 +277,57 @@ mod tests {
237277
let config = AIConfig::try_from(model).expect("conversion should succeed");
238278
assert_eq!(config.reasoning_mode, ReasoningMode::Enabled);
239279
}
280+
281+
#[test]
282+
fn derives_the_largest_output_tier_within_one_quarter_of_context() {
283+
for (context_window, expected_max_tokens) in [
284+
(32_000, 8_000),
285+
(48_000, 8_000),
286+
(64_000, 16_000),
287+
(128_000, 32_000),
288+
(128_128, 32_000),
289+
(256_000, 64_000),
290+
(1_000_000, 64_000),
291+
] {
292+
let mut model = base_model_config();
293+
model.context_window = Some(context_window);
294+
model.max_tokens = None;
295+
296+
let config = AIConfig::try_from(model).expect("conversion should succeed");
297+
298+
assert_eq!(config.max_tokens, Some(expected_max_tokens));
299+
}
300+
}
301+
302+
#[test]
303+
fn preserves_a_configured_output_limit_within_forty_percent_of_context() {
304+
let mut model = base_model_config();
305+
model.context_window = Some(1_000_000);
306+
model.max_tokens = Some(384_000);
307+
308+
let config = AIConfig::try_from(model).expect("conversion should succeed");
309+
310+
assert_eq!(config.max_tokens, Some(384_000));
311+
}
312+
313+
#[test]
314+
fn falls_back_to_the_automatic_output_limit_when_configured_limit_is_too_large() {
315+
let mut model = base_model_config();
316+
model.context_window = Some(128_000);
317+
model.max_tokens = Some(64_000);
318+
319+
let config = AIConfig::try_from(model).expect("conversion should succeed");
320+
321+
assert_eq!(config.max_tokens, Some(32_000));
322+
}
323+
324+
#[test]
325+
fn rejects_a_context_window_smaller_than_the_supported_minimum() {
326+
let mut model = base_model_config();
327+
model.context_window = Some(16_000);
328+
329+
let error = AIConfig::try_from(model).expect_err("conversion should reject small context");
330+
331+
assert!(error.contains("at least 32000"));
332+
}
240333
}

0 commit comments

Comments
 (0)