返回 CodeWhale
route_budget.rs
根目录 / crates / tui / src / route_budget.rs
1 use codewhale_config::route::RouteLimits;
2
3 use crate::config::{ApiProvider, provider_capability};
4 use crate::context_budget::ContextBudget;
5 use crate::models::{DEFAULT_COMPACTION_TOKEN_THRESHOLD, context_window_for_model};
6
7 /// Output room reserved by the internal budget for large-context reasoning
8 /// models. This is deliberately larger than the ordinary API request cap so
9 /// interleaved thinking cannot exhaust the turn budget.
10 pub(crate) const TURN_MAX_OUTPUT_TOKENS: u32 = 262_144;
11
12 /// Safe ordinary API request cap across provider routes.
13 const API_MAX_OUTPUT_TOKENS: u32 = 65_536;
14
15 /// Large windows reserve the full internal reasoning allowance. Smaller
16 /// windows reserve their route-effective request cap instead.
17 const INTERNAL_BUDGET_LARGE_WINDOW_THRESHOLD: u32 = 500_000;
18
19 /// Preserve only route limits that came from a concrete offering.
20 #[must_use]
21 pub(crate) fn known_route_limits(limits: RouteLimits) -> Option<RouteLimits> {
22 limits.has_known_limit().then_some(limits)
23 }
24
25 /// Context window for a resolved runtime route.
26 ///
27 /// Route/offering facts win when known; otherwise this falls back to the
28 /// existing provider+model capability matrix so startup and custom/local
29 /// routes keep their previous conservative behavior.
30 #[must_use]
31 pub(crate) fn route_context_window_tokens(
32 provider: ApiProvider,
33 model: &str,
34 route_limits: Option<RouteLimits>,
35 ) -> u32 {
36 route_limits
37 .and_then(|limits| limits.context_tokens)
38 .and_then(|tokens| u32::try_from(tokens).ok())
39 .filter(|tokens| *tokens > 0)
40 .unwrap_or_else(|| provider_capability(provider, model).context_window)
41 }
42
43 /// Provider/offering output cap, when the resolved route reports one.
44 #[must_use]
45 pub(crate) fn route_output_limit_tokens(route_limits: Option<RouteLimits>) -> Option<u32> {
46 route_limits
47 .and_then(|limits| limits.output_tokens)
48 .and_then(|tokens| u32::try_from(tokens).ok())
49 .filter(|tokens| *tokens > 0)
50 }
51
52 /// Effective `max_tokens` for a model before provider/route caps are applied.
53 #[must_use]
54 pub(crate) fn effective_max_output_tokens(model: &str) -> u32 {
55 if let Ok(raw) = std::env::var("CODEWHALE_MAX_OUTPUT_TOKENS")
56 .or_else(|_| std::env::var("DEEPSEEK_MAX_OUTPUT_TOKENS"))
57 && let Ok(tokens) = raw.trim().parse::<u32>()
58 && tokens > 0
59 {
60 return tokens;
61 }
62
63 let window = context_window_for_model(model).unwrap_or(128_000);
64 if window >= INTERNAL_BUDGET_LARGE_WINDOW_THRESHOLD {
65 API_MAX_OUTPUT_TOKENS
66 } else {
67 (window / 2).min(API_MAX_OUTPUT_TOKENS)
68 }
69 }
70
71 /// Conservative request ceiling for a model the static catalogue does not
72 /// describe at all.
73 ///
74 /// An absent compatibility cap is not evidence of a large ceiling. Remote
75 /// OpenAI-compatible routes serving an unrecognized wire alias frequently
76 /// publish a much lower `max_tokens` maximum and reject anything above it, so
77 /// an uncatalogued id keeps this floor rather than inheriting the full
78 /// [`API_MAX_OUTPUT_TOKENS`] request cap.
79 const UNCATALOGUED_COMPAT_MAX_OUTPUT_TOKENS: u32 = 8_192;
80
81 /// Why a route's compatibility output ceiling has the value it does.
82 ///
83 /// Carried so a clamp is always attributable: "unknown" is only allowed to
84 /// mean "no clamp" when a route *truthfully publishes no ceiling*, never when
85 /// the catalogue simply has no row for the model.
86 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
87 pub(crate) enum OutputCeilingSource {
88 /// The static catalogue publishes an exact/conservative ceiling.
89 Documented(u32),
90 /// The route is known to publish no output maximum we can stand behind
91 /// (Kimi Code membership ids, operator-owned self-hosted engines). Unknown
92 /// stays unknown and nothing is clamped.
93 RouteDeclaredUnknown,
94 /// The catalogue has no row for this model. Fail closed to a conservative
95 /// ceiling rather than treating absence as permission.
96 Uncatalogued(u32),
97 }
98
99 impl OutputCeilingSource {
100 /// The ceiling to intersect a requested cap with, if any.
101 #[must_use]
102 pub(crate) const fn clamp_tokens(self) -> Option<u32> {
103 match self {
104 Self::Documented(tokens) | Self::Uncatalogued(tokens) => Some(tokens),
105 Self::RouteDeclaredUnknown => None,
106 }
107 }
108 }
109
110 /// Whether an absent compatibility ceiling is a *declared* unknown for this
111 /// route, rather than a gap in the catalogue.
112 ///
113 /// Deliberately an allowlist. Everything not named here is uncatalogued and
114 /// gets the conservative ceiling.
115 #[must_use]
116 fn route_declares_unknown_output_ceiling(provider: ApiProvider, model: &str) -> bool {
117 match provider {
118 // Operator-owned engines: the local server, not this process, owns the
119 // output ceiling, and it is routinely far above any catalogue row.
120 ApiProvider::Ollama | ApiProvider::Sglang | ApiProvider::Vllm => true,
121 // Kimi Code membership ids publish their limits in the membership
122 // catalog rather than the static model catalogue.
123 ApiProvider::Moonshot => crate::config::is_kimi_code_membership_model(model),
124 _ => false,
125 }
126 }
127
128 /// Resolve the compatibility output ceiling for a route, with its provenance.
129 #[must_use]
130 pub(crate) fn output_ceiling_source(provider: ApiProvider, model: &str) -> OutputCeilingSource {
131 provider_capability(provider, model).max_output.map_or_else(
132 || {
133 if route_declares_unknown_output_ceiling(provider, model) {
134 OutputCeilingSource::RouteDeclaredUnknown
135 } else {
136 OutputCeilingSource::Uncatalogued(UNCATALOGUED_COMPAT_MAX_OUTPUT_TOKENS)
137 }
138 },
139 OutputCeilingSource::Documented,
140 )
141 }
142
143 /// Effective request output cap for a fully resolved provider/model route.
144 #[must_use]
145 pub(crate) fn effective_max_output_tokens_for_route(
146 provider: ApiProvider,
147 model: &str,
148 route_limits: Option<RouteLimits>,
149 ) -> u32 {
150 let requested_cap = effective_max_output_tokens(model);
151 let compatibility_cap = output_ceiling_source(provider, model).clamp_tokens();
152 let route_cap = route_output_limit_tokens(route_limits);
153 // Unknown means unknown only where a route *declares* it: membership ids
154 // such as the `kimi-for-coding` family, and operator-owned self-hosted
155 // engines. For those there is nothing to clamp against and the requested
156 // cap stands. A model the catalogue simply has no row for is not the same
157 // fact — absence is not permission, so it keeps a conservative ceiling
158 // (see `output_ceiling_source`). Only a concrete route/offering maximum
159 // narrows it further; known compatibility caps stay authoritative and are
160 // still intersected with any route maximum.
161 let cap = compatibility_cap.map_or(requested_cap, |compat| requested_cap.min(compat));
162 let cap = route_cap.map_or(cap, |route_cap| cap.min(route_cap));
163 let Some(window) = route_limits
164 .and_then(|limits| limits.context_tokens)
165 .and_then(|tokens| u32::try_from(tokens).ok())
166 .filter(|tokens| *tokens > 0)
167 else {
168 return cap;
169 };
170
171 u32::try_from(ContextBudget::new(u64::from(window), 0, u64::from(cap)).output_cap_tokens)
172 .unwrap_or(cap)
173 .max(1)
174 }
175
176 /// Output reservation used by the internal input budget for a route.
177 #[must_use]
178 pub(crate) fn route_output_reservation_for_window(
179 provider: ApiProvider,
180 model: &str,
181 window_tokens: u32,
182 route_limits: Option<RouteLimits>,
183 ) -> u32 {
184 if window_tokens >= INTERNAL_BUDGET_LARGE_WINDOW_THRESHOLD {
185 route_output_limit_tokens(route_limits).map_or(TURN_MAX_OUTPUT_TOKENS, |route_cap| {
186 route_cap.min(TURN_MAX_OUTPUT_TOKENS)
187 })
188 } else {
189 effective_max_output_tokens_for_route(provider, model, route_limits)
190 }
191 }
192
193 #[must_use]
194 pub(crate) fn route_context_budget(
195 provider: ApiProvider,
196 model: &str,
197 route_limits: Option<RouteLimits>,
198 input_tokens: usize,
199 ) -> Option<ContextBudget> {
200 let window = route_context_window_tokens(provider, model, route_limits);
201 let output_cap = route_output_reservation_for_window(provider, model, window, route_limits);
202 Some(ContextBudget::new(
203 u64::from(window),
204 u64::try_from(input_tokens).ok()?,
205 u64::from(output_cap),
206 ))
207 }
208
209 #[must_use]
210 pub(crate) fn compaction_threshold_for_route_at_percent(
211 provider: ApiProvider,
212 model: &str,
213 route_limits: Option<RouteLimits>,
214 percent: f64,
215 ) -> usize {
216 route_context_budget(provider, model, route_limits, 0)
217 .and_then(|budget| {
218 usize::try_from(budget.compaction_trigger_for_percent(percent.clamp(10.0, 100.0))).ok()
219 })
220 .unwrap_or(DEFAULT_COMPACTION_TOKEN_THRESHOLD)
221 }
222
223 #[must_use]
224 pub(crate) fn auto_compact_default_for_route(
225 provider: ApiProvider,
226 model: &str,
227 route_limits: Option<RouteLimits>,
228 ) -> bool {
229 // Every resolved route has either concrete offering limits or a
230 // conservative provider/model fallback. Large windows need continuity too;
231 // their size is not a reason to disable compaction entirely.
232 route_context_window_tokens(provider, model, route_limits) > 0
233 }
234
235 #[cfg(test)]
236 mod tests {
237 use super::*;
238
239 /// Absence of a catalogue row is not evidence of a large ceiling. An
240 /// unrecognized wire alias on a remote OpenAI-compatible route keeps the
241 /// conservative compatibility ceiling, with an attributable source.
242 #[test]
243 fn uncatalogued_remote_model_keeps_a_conservative_ceiling() {
244 let source = output_ceiling_source(ApiProvider::Openai, "totally-unknown-alias-v9");
245 assert_eq!(
246 source,
247 OutputCeilingSource::Uncatalogued(UNCATALOGUED_COMPAT_MAX_OUTPUT_TOKENS)
248 );
249 assert_eq!(
250 source.clamp_tokens(),
251 Some(UNCATALOGUED_COMPAT_MAX_OUTPUT_TOKENS)
252 );
253 assert!(
254 effective_max_output_tokens_for_route(
255 ApiProvider::Openai,
256 "totally-unknown-alias-v9",
257 None
258 ) <= UNCATALOGUED_COMPAT_MAX_OUTPUT_TOKENS
259 );
260 }
261
262 /// Routes that *declare* an unknown ceiling still avoid the clamp.
263 #[test]
264 fn route_declared_unknown_ceilings_are_not_clamped() {
265 for (provider, model) in [
266 (ApiProvider::Moonshot, "kimi-for-coding"),
267 (ApiProvider::Moonshot, "kimi-for-coding-highspeed"),
268 (ApiProvider::Ollama, "some-local-build"),
269 ] {
270 assert_eq!(
271 output_ceiling_source(provider, model),
272 OutputCeilingSource::RouteDeclaredUnknown,
273 "{provider:?}/{model} must declare its unknown ceiling"
274 );
275 assert_eq!(output_ceiling_source(provider, model).clamp_tokens(), None);
276 }
277 // Bare `k3` is a membership id, but unlike the `kimi-for-coding`
278 // family the K3 quickstart documents its output maximum, and the model
279 // catalogue carries it. A documented ceiling is authoritative — the
280 // membership allowlist only covers ids the catalogue has nothing to
281 // say about, and must not turn a real fact back into an unknown.
282 assert_eq!(
283 output_ceiling_source(ApiProvider::Moonshot, "k3"),
284 OutputCeilingSource::Documented(131_072)
285 );
286 }
287
288 #[test]
289 fn codex_missing_route_metadata_uses_provider_context_floor() {
290 assert_eq!(
291 route_context_window_tokens(ApiProvider::OpenaiCodex, "gpt-5.5", None),
292 128_000
293 );
294 assert_eq!(
295 compaction_threshold_for_route_at_percent(
296 ApiProvider::OpenaiCodex,
297 "gpt-5.5",
298 None,
299 80.0,
300 ),
301 98_304
302 );
303 assert!(auto_compact_default_for_route(
304 ApiProvider::OpenaiCodex,
305 "gpt-5.5",
306 None,
307 ));
308 }
309
310 #[test]
311 fn v4_trigger_is_anchored_to_spendable_input() {
312 let budget = route_context_budget(ApiProvider::Deepseek, "deepseek-v4-pro", None, 0)
313 .expect("V4 route budget");
314
315 assert_eq!(budget.window_tokens, 1_000_000);
316 assert_eq!(budget.output_cap_tokens, u64::from(TURN_MAX_OUTPUT_TOKENS));
317 assert_eq!(budget.input_budget_ceiling, 736_832);
318 assert_eq!(
319 compaction_threshold_for_route_at_percent(
320 ApiProvider::Deepseek,
321 "deepseek-v4-pro",
322 None,
323 80.0,
324 ),
325 589_466
326 );
327 }
328
329 #[test]
330 fn kimi_k3_defaults_auto_compaction_on() {
331 assert!(auto_compact_default_for_route(
332 ApiProvider::Moonshot,
333 "kimi-k3",
334 None,
335 ));
336 }
337
338 #[test]
339 fn kimi_catalog_output_ceiling_preserves_input_budget() {
340 let _lock = crate::test_support::lock_test_env();
341 let _max_output = crate::test_support::EnvVarGuard::remove("DEEPSEEK_MAX_OUTPUT_TOKENS");
342 // #4368/#4378: Models.dev may report Kimi's full 262K context as both
343 // context and output ceilings. On a sub-500K window, reserve the
344 // route-effective 32K request cap rather than treating that catalog
345 // maximum as the amount every turn will emit.
346 let limits = RouteLimits {
347 context_tokens: Some(262_144),
348 output_tokens: Some(262_144),
349 ..RouteLimits::default()
350 };
351 let budget = route_context_budget(ApiProvider::Moonshot, "kimi-k2.7-code", Some(limits), 0)
352 .expect("Kimi route budget");
353 let trigger = compaction_threshold_for_route_at_percent(
354 ApiProvider::Moonshot,
355 "kimi-k2.7-code",
356 Some(limits),
357 80.0,
358 );
359
360 assert_eq!(budget.output_cap_tokens, 32_768);
361 assert_eq!(budget.input_budget_ceiling, 228_352);
362 assert_eq!(trigger, 182_682);
363 assert!(trigger as u64 <= budget.input_budget_ceiling);
364 assert!(
365 trigger < 209_715,
366 "must fire before the old window-relative trigger"
367 );
368 }
369
370 #[test]
371 fn explicit_route_output_limit_beats_unknown_model_name_fallback() {
372 let _lock = crate::test_support::lock_test_env();
373 let _max_output =
374 crate::test_support::EnvVarGuard::set("CODEWHALE_MAX_OUTPUT_TOKENS", "65536");
375 let limits = RouteLimits {
376 context_tokens: Some(262_144),
377 output_tokens: Some(24_576),
378 ..RouteLimits::default()
379 };
380
381 assert_eq!(
382 effective_max_output_tokens_for_route(
383 ApiProvider::Vllm,
384 "arbitrary-local-wire-alias",
385 Some(limits),
386 ),
387 24_576
388 );
389 assert_eq!(
390 effective_max_output_tokens_for_route(
391 ApiProvider::Vllm,
392 "arbitrary-local-wire-alias",
393 None,
394 ),
395 65_536,
396 "an unknown compatibility cap must not clamp; only the requested cap applies"
397 );
398 assert_eq!(
399 effective_max_output_tokens_for_route(
400 ApiProvider::Vllm,
401 "kimi-k2.7-code",
402 Some(RouteLimits {
403 output_tokens: Some(262_144),
404 ..RouteLimits::default()
405 }),
406 ),
407 32_768,
408 "known model caps must remain authoritative on self-hosted routes"
409 );
410 }
411
412 /// #4368 follow-up: the Kimi Code membership ids deliberately have no
413 /// static output cap (the membership catalog owns their limits). The old
414 /// generic `unwrap_or(4096)` in `provider_capability` turned that unknown
415 /// into a hard 4K clamp here, silently truncating every offline membership
416 /// turn. Unknown must mean "no compatibility clamp".
417 #[test]
418 fn kimi_membership_unknown_output_cap_does_not_clamp_to_4k() {
419 let _lock = crate::test_support::lock_test_env();
420 let _codewhale = crate::test_support::EnvVarGuard::remove("CODEWHALE_MAX_OUTPUT_TOKENS");
421 let _deepseek = crate::test_support::EnvVarGuard::remove("DEEPSEEK_MAX_OUTPUT_TOKENS");
422
423 for model in ["kimi-for-coding", "kimi-for-coding-highspeed"] {
424 assert_eq!(
425 provider_capability(ApiProvider::Moonshot, model).max_output,
426 None,
427 "{model}: membership output ceiling must stay unknown, not a placeholder"
428 );
429
430 let cap = effective_max_output_tokens_for_route(ApiProvider::Moonshot, model, None);
431 assert_eq!(
432 cap,
433 effective_max_output_tokens(model),
434 "{model}: unknown compatibility cap must leave the requested cap intact"
435 );
436 assert_ne!(cap, 4_096, "{model}: must not inherit the old 4K fallback");
437 // No invented sentinel ceiling either.
438 assert_ne!(cap, u32::MAX);
439 assert_ne!(cap, 32_768);
440 }
441 }
442
443 /// A concrete membership offering limit is still authoritative — "unknown
444 /// means no clamp" must not become "never clamp".
445 #[test]
446 fn kimi_membership_route_limit_still_caps_output() {
447 let _lock = crate::test_support::lock_test_env();
448 let _codewhale = crate::test_support::EnvVarGuard::remove("CODEWHALE_MAX_OUTPUT_TOKENS");
449 let _deepseek = crate::test_support::EnvVarGuard::remove("DEEPSEEK_MAX_OUTPUT_TOKENS");
450
451 let limits = RouteLimits {
452 context_tokens: Some(262_144),
453 output_tokens: Some(16_384),
454 ..RouteLimits::default()
455 };
456 assert_eq!(
457 effective_max_output_tokens_for_route(
458 ApiProvider::Moonshot,
459 "kimi-for-coding",
460 Some(limits),
461 ),
462 16_384
463 );
464 }
465
466 /// GLM and MiniMax publish real output ceilings; those stay authoritative
467 /// so relaxing the unknown case cannot leak into known routes.
468 #[test]
469 fn known_glm_and_minimax_output_caps_remain_authoritative() {
470 let _lock = crate::test_support::lock_test_env();
471 let _codewhale = crate::test_support::EnvVarGuard::remove("CODEWHALE_MAX_OUTPUT_TOKENS");
472 let _deepseek = crate::test_support::EnvVarGuard::remove("DEEPSEEK_MAX_OUTPUT_TOKENS");
473
474 // GLM 5.2: 1M window, documented 131K output. The requested cap is the
475 // 65,536 API ceiling, so the known cap is above it and does not bind —
476 // what matters is that the capability is *known*.
477 let glm = provider_capability(ApiProvider::Zai, "glm-5.2");
478 assert_eq!(glm.max_output, Some(131_072));
479
480 let minimax = provider_capability(ApiProvider::Minimax, "minimax-m3");
481 assert_eq!(minimax.max_output, Some(524_288));
482
483 // A known cap below the requested cap must still clamp.
484 assert_eq!(
485 effective_max_output_tokens_for_route(ApiProvider::Moonshot, "kimi-k2.7-code", None),
486 32_768,
487 );
488 }
489 }
490
490 lines RUST