1use std::{
8 collections::{HashMap, HashSet},
9 str,
10 sync::Arc,
11};
12
13use ferrum_interfaces::tokenizer::Tokenizer;
14use ferrum_types::{FerrumError, ResponseFormat, Result, StructuredOutputStart, TokenId};
15use llguidance::{
16 api::TopLevelGrammar,
17 toktrie::{InferenceCapabilities, TokEnv, TokRxInfo, TokTrie, TokenizerEnv},
18 JsonCompileOptions, Matcher, ParserFactory,
19};
20use parking_lot::Mutex;
21use serde_json::json;
22
23mod auto_tools;
24mod harmony;
25#[cfg(test)]
26mod reasoning_envelope_tests;
27
28const MAX_CACHED_GRAMMARS: usize = 64;
29const AUTO_STRUCTURED_RESERVE_DIVISOR: usize = 2;
32const MIN_AUTO_STRUCTURED_RESERVE_TOKENS: usize = 32;
33const MAX_AUTO_STRUCTURED_RESERVE_TOKENS: usize = 1024;
34const MAX_IDENTICAL_TOKEN_RUN: usize = MAX_AUTO_STRUCTURED_RESERVE_TOKENS / 2;
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct StructuredOutputBudgetPlan {
40 pub total_output_tokens: usize,
41 pub reasoning_token_limit: usize,
42 pub boundary_token_count: usize,
43 pub structured_reserve_tokens: usize,
44}
45
46impl StructuredOutputBudgetPlan {
47 fn automatic(total_output_tokens: usize, boundary_token_count: usize) -> Result<Self> {
48 if boundary_token_count == 0 || total_output_tokens <= boundary_token_count {
49 return Err(FerrumError::invalid_request(format!(
50 "structured output requires max_tokens greater than its {boundary_token_count}-token delimiter"
51 )));
52 }
53 let available_after_boundary = total_output_tokens - boundary_token_count;
54 let proportional_reserve = total_output_tokens.div_ceil(AUTO_STRUCTURED_RESERVE_DIVISOR);
55 let structured_reserve_tokens = proportional_reserve
56 .clamp(
57 MIN_AUTO_STRUCTURED_RESERVE_TOKENS,
58 MAX_AUTO_STRUCTURED_RESERVE_TOKENS,
59 )
60 .min(available_after_boundary);
61 Ok(Self {
62 total_output_tokens,
63 reasoning_token_limit: total_output_tokens
64 - boundary_token_count
65 - structured_reserve_tokens,
66 boundary_token_count,
67 structured_reserve_tokens,
68 })
69 }
70}
71
72#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73struct StructuredOutputLivenessPolicy {
74 max_identical_token_run: usize,
75}
76
77impl StructuredOutputLivenessPolicy {
78 fn for_request(max_output_tokens: usize, budget: Option<StructuredOutputBudgetPlan>) -> Self {
79 let guaranteed_structured_tokens = budget
80 .map(|plan| plan.structured_reserve_tokens)
81 .unwrap_or(max_output_tokens);
82 Self {
83 max_identical_token_run: guaranteed_structured_tokens
84 .div_ceil(2)
85 .clamp(1, MAX_IDENTICAL_TOKEN_RUN),
86 }
87 }
88}
89
90pub struct StructuredOutputFactory {
92 parser_factory: ParserFactory,
93 tokenizer: Arc<dyn Tokenizer + Send + Sync>,
94 vocab_size: usize,
95 defined_token_ids: Arc<[bool]>,
96 json_token_classes: Arc<[StructuredOutputTokenClass]>,
97 wire_token_bytes: Arc<[Vec<u8>]>,
98 grammar_templates: Mutex<HashMap<String, Matcher>>,
99 schema_validators: Mutex<HashMap<String, Arc<jsonschema::Validator>>>,
100}
101
102impl std::fmt::Debug for StructuredOutputFactory {
103 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104 f.debug_struct("StructuredOutputFactory")
105 .field("vocab_size", &self.vocab_size)
106 .finish_non_exhaustive()
107 }
108}
109
110impl StructuredOutputFactory {
111 pub fn new(tokenizer: Arc<dyn Tokenizer + Send + Sync>) -> Result<Self> {
113 Self::new_with_model_vocab_size(tokenizer, None)
114 }
115
116 pub fn new_with_model_vocab_size(
119 tokenizer: Arc<dyn Tokenizer + Send + Sync>,
120 model_vocab_size: Option<usize>,
121 ) -> Result<Self> {
122 let eos = tokenizer.special_tokens().eos_token.ok_or_else(|| {
123 FerrumError::config("structured output requires a tokenizer EOS token")
124 })?;
125 let vocab_size = model_vocab_size
126 .unwrap_or_else(|| tokenizer.vocab_size())
127 .max(tokenizer.vocab_size());
128 if vocab_size == 0 || eos.get() as usize >= vocab_size {
129 return Err(FerrumError::config(format!(
130 "structured output tokenizer has invalid vocab/EOS: vocab_size={vocab_size}, eos={}",
131 eos.get()
132 )));
133 }
134
135 let special_ids = tokenizer_special_ids(tokenizer.as_ref());
136 let wire_token_bytes = (0..vocab_size)
137 .map(|idx| {
138 tokenizer
139 .token_bytes(TokenId::new(idx as u32))
140 .unwrap_or_default()
141 })
142 .collect::<Vec<_>>();
143 let mut defined_token_ids = Vec::with_capacity(vocab_size);
144 let mut json_token_classes = Vec::with_capacity(vocab_size);
145 let token_bytes = (0..vocab_size)
146 .map(|idx| {
147 let token = TokenId::new(idx as u32);
148 if special_ids.contains(&token.get()) {
149 defined_token_ids.push(true);
150 json_token_classes.push(StructuredOutputTokenClass::Control);
151 special_token_marker(token)
152 } else if let Some(bytes) = tokenizer
153 .token_bytes(token)
154 .filter(|bytes| !bytes.is_empty())
155 {
156 defined_token_ids.push(true);
157 json_token_classes.push(classify_json_token_bytes(&bytes));
158 bytes
159 } else {
160 defined_token_ids.push(false);
165 json_token_classes.push(StructuredOutputTokenClass::Undefined);
166 Vec::new()
167 }
168 })
169 .collect::<Vec<_>>();
170
171 let mut eos_tokens = vec![eos.get()];
172 eos_tokens.extend(
173 tokenizer
174 .special_tokens()
175 .extra_eos_tokens
176 .iter()
177 .map(|token| token.get())
178 .filter(|token| *token < vocab_size as u32),
179 );
180 eos_tokens.sort_unstable();
181 eos_tokens.dedup();
182 if let Some(position) = eos_tokens.iter().position(|token| *token == eos.get()) {
183 eos_tokens.swap(0, position);
184 }
185
186 let info = TokRxInfo::new(vocab_size as u32, eos.get());
187 let trie = TokTrie::from(&info, &token_bytes).with_eos_tokens(&eos_tokens);
188 let tok_env: TokEnv = Arc::new(FerrumTokenizerEnv {
189 tokenizer: Arc::clone(&tokenizer),
190 trie,
191 });
192 let mut parser_factory = ParserFactory::new(
193 &tok_env,
194 InferenceCapabilities {
195 ff_tokens: false,
196 conditional_ff_tokens: false,
197 backtrack: false,
198 fork: false,
199 },
200 &llguidance::earley::SlicedBiasComputer::general_slices(),
201 )
202 .map_err(|error| {
203 FerrumError::config(format!("build structured-output parser factory: {error}"))
204 })?;
205 parser_factory.quiet();
206
207 Ok(Self {
208 parser_factory,
209 tokenizer,
210 vocab_size,
211 defined_token_ids: defined_token_ids.into(),
212 json_token_classes: json_token_classes.into(),
213 wire_token_bytes: wire_token_bytes.into(),
214 grammar_templates: Mutex::new(HashMap::new()),
215 schema_validators: Mutex::new(HashMap::new()),
216 })
217 }
218
219 pub fn create_processor_with_chat_contract(
221 &self,
222 response_format: &ResponseFormat,
223 start: &StructuredOutputStart,
224 max_output_tokens: usize,
225 stop_token_ids: &HashSet<u32>,
226 stop_text_sequences: &[String],
227 chat_request: Option<&ferrum_types::ApiChatRequest>,
228 output_protocol: ferrum_types::ModelOutputProtocol,
229 ) -> Result<Option<StructuredOutputProcessor>> {
230 if let Some(chat) = chat_request.filter(|chat| {
231 chat.automatic_tools_with_hard_response_format()
232 && !matches!(response_format, ResponseFormat::Text)
233 }) {
234 return auto_tools::compile(
235 self,
236 response_format,
237 start,
238 max_output_tokens,
239 stop_token_ids,
240 stop_text_sequences,
241 chat,
242 output_protocol,
243 )
244 .map(Some);
245 }
246 self.create_processor(
247 response_format,
248 start,
249 max_output_tokens,
250 stop_token_ids,
251 stop_text_sequences,
252 )
253 }
254
255 pub fn create_processor(
257 &self,
258 response_format: &ResponseFormat,
259 start: &StructuredOutputStart,
260 max_output_tokens: usize,
261 stop_token_ids: &HashSet<u32>,
262 stop_text_sequences: &[String],
263 ) -> Result<Option<StructuredOutputProcessor>> {
264 let schema = match response_format {
265 ResponseFormat::Text => return Ok(None),
266 ResponseFormat::JsonObject => json!({"type": "object"}),
267 ResponseFormat::JsonSchema(schema) => {
268 serde_json::from_str(schema).map_err(|error| {
269 FerrumError::invalid_request(format!(
270 "response_format.schema is not valid JSON: {error}"
271 ))
272 })?
273 }
274 };
275 let schema = compact_json_schema(schema)?;
276 let grammar_key = serde_json::to_string(&schema).map_err(|error| {
277 FerrumError::invalid_request(format!("serialize structured-output schema: {error}"))
278 })?;
279 let matcher = {
280 let mut templates = self.grammar_templates.lock();
281 if let Some(template) = templates.get(&grammar_key) {
282 template.deep_clone()
283 } else {
284 let grammar = TopLevelGrammar::from_json_schema(schema);
285 let parser = self
286 .parser_factory
287 .create_parser(grammar)
288 .map_err(|error| {
289 FerrumError::invalid_request(format!(
290 "unsupported structured-output grammar: {error}"
291 ))
292 })?;
293 let matcher = Matcher::new(Ok(parser));
294 if templates.len() >= MAX_CACHED_GRAMMARS {
295 templates.clear();
296 }
297 templates.insert(grammar_key, matcher.deep_clone());
298 matcher
299 }
300 };
301 let (activation, budget) = match start {
302 StructuredOutputStart::Immediate => (Activation::Active, None),
303 StructuredOutputStart::HarmonyFinal => {
304 let boundary = harmony::HarmonyBoundary::compile(
305 self.tokenizer.as_ref(),
306 self.vocab_size,
307 max_output_tokens,
308 stop_token_ids,
309 stop_text_sequences,
310 )?;
311 let budget = boundary.budget();
312 (Activation::Harmony(boundary), Some(budget))
313 }
314 StructuredOutputStart::AfterDelimiter(delimiter) => {
315 let delimiter_tokens =
316 self.boundary_token_ids(delimiter, stop_token_ids, stop_text_sequences)?;
317 let budget = StructuredOutputBudgetPlan::automatic(
318 max_output_tokens,
319 delimiter_tokens.len(),
320 )?;
321 (
322 Activation::Boundary {
323 opening_tokens: Vec::new(),
324 delimiter_tokens,
325 allow_reasoning: true,
326 forcing: false,
327 },
328 Some(budget),
329 )
330 }
331 StructuredOutputStart::AfterReasoningEnvelope {
332 opening,
333 closing,
334 allow_reasoning,
335 } => {
336 let opening_tokens =
337 self.boundary_token_ids(opening, stop_token_ids, stop_text_sequences)?;
338 let delimiter_tokens =
339 self.boundary_token_ids(closing, stop_token_ids, stop_text_sequences)?;
340 if let Some(token) = opening_tokens
341 .iter()
342 .chain(&delimiter_tokens)
343 .find(|token| {
344 !self
345 .defined_token_ids
346 .get(**token as usize)
347 .copied()
348 .unwrap_or(false)
349 })
350 {
351 return Err(FerrumError::invalid_request(format!(
352 "structured-output reasoning envelope contains undefined token {token}"
353 )));
354 }
355 let empty_envelope = format!("{opening}{closing}");
358 if let Some(stop) = stop_text_sequences
359 .iter()
360 .find(|stop| !stop.is_empty() && empty_envelope.contains(stop.as_str()))
361 {
362 return Err(FerrumError::invalid_request(format!(
363 "structured-output reasoning envelope conflicts with stop sequence {stop:?}"
364 )));
365 }
366 let mut budget = StructuredOutputBudgetPlan::automatic(
367 max_output_tokens,
368 opening_tokens.len() + delimiter_tokens.len(),
369 )?;
370 if !allow_reasoning {
371 budget.reasoning_token_limit = 0;
372 budget.structured_reserve_tokens =
373 max_output_tokens - budget.boundary_token_count;
374 }
375 (
376 Activation::Boundary {
377 opening_tokens,
378 delimiter_tokens,
379 allow_reasoning: *allow_reasoning,
380 forcing: false,
381 },
382 Some(budget),
383 )
384 }
385 };
386
387 let grammar_start = matches!(activation, Activation::Active).then_some(0);
388 let liveness = StructuredOutputLivenessPolicy::for_request(max_output_tokens, budget);
389 Ok(Some(StructuredOutputProcessor {
390 state: Mutex::new(ProcessorState {
391 matcher,
392 activation: activation.clone(),
393 initial_activation: activation,
394 consumed: 0,
395 boundary_forced: false,
396 boundary_start: None,
397 grammar_start,
398 trailing_grammar_token_id: None,
399 trailing_identical_token_count: 0,
400 liveness_intervention_count: 0,
401 last_liveness_intervention_at: None,
402 composed: None,
403 }),
404 vocab_size: self.vocab_size,
405 defined_token_ids: Arc::clone(&self.defined_token_ids),
406 json_token_classes: Arc::clone(&self.json_token_classes),
407 budget,
408 liveness,
409 }))
410 }
411
412 fn boundary_token_ids(
413 &self,
414 delimiter: &str,
415 stop_token_ids: &HashSet<u32>,
416 stop_text_sequences: &[String],
417 ) -> Result<Vec<u32>> {
418 if delimiter.is_empty() {
419 return Err(FerrumError::invalid_request(
420 "structured-output delimiter must not be empty",
421 ));
422 }
423 let tokens = if let Some(token) = self.tokenizer.token_id(delimiter) {
424 vec![token.get()]
425 } else {
426 self.tokenizer
427 .encode(delimiter, false)?
428 .into_iter()
429 .map(|token| token.get())
430 .collect::<Vec<_>>()
431 };
432 if tokens.is_empty() {
433 return Err(FerrumError::invalid_request(format!(
434 "structured-output delimiter {delimiter:?} did not tokenize"
435 )));
436 }
437 if let Some(token) = tokens.iter().find(|token| stop_token_ids.contains(token)) {
438 return Err(FerrumError::invalid_request(format!(
439 "structured-output delimiter token {token} conflicts with a stop token"
440 )));
441 }
442 if let Some(stop) = stop_text_sequences
443 .iter()
444 .find(|stop| !stop.is_empty() && delimiter.contains(stop.as_str()))
445 {
446 return Err(FerrumError::invalid_request(format!(
447 "structured-output delimiter {delimiter:?} conflicts with stop sequence {stop:?}"
448 )));
449 }
450 Ok(tokens)
451 }
452}
453
454fn compact_json_schema(schema: serde_json::Value) -> Result<serde_json::Value> {
455 let mut schema = match schema {
456 schema @ serde_json::Value::Object(_) => schema,
457 schema @ serde_json::Value::Bool(_) => json!({"allOf": [schema]}),
458 _ => {
459 return Err(FerrumError::invalid_request(
460 "response_format.schema must be a JSON Schema object or boolean",
461 ));
462 }
463 };
464
465 JsonCompileOptions {
469 whitespace_flexible: false,
470 ..JsonCompileOptions::default()
471 }
472 .apply_to(&mut schema);
473 Ok(schema)
474}
475
476pub struct StructuredOutputProcessor {
478 state: Mutex<ProcessorState>,
479 vocab_size: usize,
480 defined_token_ids: Arc<[bool]>,
481 json_token_classes: Arc<[StructuredOutputTokenClass]>,
482 budget: Option<StructuredOutputBudgetPlan>,
483 liveness: StructuredOutputLivenessPolicy,
484}
485
486#[derive(Debug, Clone, Copy, PartialEq, Eq)]
488pub enum StructuredOutputPhase {
489 WaitingForDelimiter,
490 ForcingDelimiter,
491 EnforcingGrammar,
492}
493
494#[derive(Debug, Clone, Copy, PartialEq, Eq)]
500#[repr(u8)]
501pub enum StructuredOutputTokenClass {
502 Whitespace,
503 Number,
504 Structural,
505 StringBoundary,
506 Literal,
507 Other,
508 Control,
509 Undefined,
510}
511
512#[derive(Debug, Clone, Copy, PartialEq, Eq)]
514pub struct StructuredOutputMaskOutcome {
515 pub phase: StructuredOutputPhase,
516 pub accepting: bool,
517 pub liveness_intervention: bool,
518 pub grammar_start_token_index: Option<usize>,
523 pub required_delimiter_token_id: Option<u32>,
528 pub grammar_owned_terminal_token_id: Option<u32>,
531}
532
533#[derive(Debug, Clone, Copy, PartialEq, Eq)]
536pub struct StructuredOutputProgress {
537 pub phase: StructuredOutputPhase,
538 pub generated_token_count: usize,
539 pub consumed_token_count: usize,
540 pub delimiter_token_count: Option<usize>,
541 pub delimiter_prefix_token_count: usize,
542 pub reasoning_token_count: Option<usize>,
543 pub boundary_forced: bool,
544 pub budget: Option<StructuredOutputBudgetPlan>,
545 pub grammar_token_count: usize,
546 pub trailing_token_class: Option<StructuredOutputTokenClass>,
547 pub trailing_token_class_count: usize,
548 pub trailing_token_id: Option<u32>,
549 pub trailing_identical_token_count: usize,
550 pub liveness_identical_token_limit: usize,
551 pub liveness_intervention_count: usize,
552 pub accepting: bool,
553}
554
555impl std::fmt::Debug for StructuredOutputProcessor {
556 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
557 let state = self.state.lock();
558 f.debug_struct("StructuredOutputProcessor")
559 .field("vocab_size", &self.vocab_size)
560 .field("consumed", &state.consumed)
561 .field("active", &state.grammar_start.is_some())
562 .field("budget", &self.budget)
563 .finish()
564 }
565}
566
567struct ProcessorState {
568 matcher: Matcher,
569 activation: Activation,
570 initial_activation: Activation,
571 consumed: usize,
572 boundary_forced: bool,
573 boundary_start: Option<usize>,
574 grammar_start: Option<usize>,
575 trailing_grammar_token_id: Option<u32>,
576 trailing_identical_token_count: usize,
577 liveness_intervention_count: usize,
578 last_liveness_intervention_at: Option<usize>,
579 composed: Option<auto_tools::ComposedState>,
580}
581
582#[derive(Clone)]
583enum Activation {
584 Active,
585 Harmony(harmony::HarmonyBoundary),
586 Boundary {
587 opening_tokens: Vec<u32>,
588 delimiter_tokens: Vec<u32>,
589 allow_reasoning: bool,
590 forcing: bool,
591 },
592}
593
594impl StructuredOutputProcessor {
595 pub fn mask_logits(&self, logits: &mut [f32], generated: &[TokenId]) -> Result<()> {
599 self.mask_logits_inner(logits, generated, None, None)
600 .map(|_| ())
601 }
602
603 pub fn mask_logits_with_terminals(
608 &self,
609 logits: &mut [f32],
610 generated: &[TokenId],
611 terminal_token_ids: &HashSet<u32>,
612 hidden_control_token_ids: &HashSet<u32>,
613 ) -> Result<StructuredOutputMaskOutcome> {
614 self.mask_logits_inner(
615 logits,
616 generated,
617 Some(terminal_token_ids),
618 Some(hidden_control_token_ids),
619 )
620 }
621
622 fn mask_logits_inner(
623 &self,
624 logits: &mut [f32],
625 generated: &[TokenId],
626 terminal_token_ids: Option<&HashSet<u32>>,
627 hidden_control_token_ids: Option<&HashSet<u32>>,
628 ) -> Result<StructuredOutputMaskOutcome> {
629 let mut state = self.state.lock();
630 if state.composed.is_some() {
631 return auto_tools::mask(
632 self,
633 &mut state,
634 logits,
635 generated,
636 terminal_token_ids,
637 hidden_control_token_ids,
638 );
639 }
640 advance_state(&mut state, generated, terminal_token_ids, self.budget)?;
641 activate_forcing_if_due(&mut state, generated, self.budget);
642 if let Activation::Harmony(boundary) = &state.activation {
643 if !boundary.in_payload() {
644 self.mask_undefined_token_ids(logits);
645 let required_delimiter_token_id =
646 boundary.mask_before_payload(logits, hidden_control_token_ids)?;
647 return Ok(StructuredOutputMaskOutcome {
648 phase: boundary.progress().0,
649 accepting: false,
650 liveness_intervention: false,
651 grammar_start_token_index: None,
652 required_delimiter_token_id,
653 grammar_owned_terminal_token_id: None,
654 });
655 }
656 }
657 if let Activation::Boundary {
658 opening_tokens,
659 delimiter_tokens,
660 forcing,
661 ..
662 } = &state.activation
663 {
664 self.mask_undefined_token_ids(logits);
665 let opening_token = opening_tokens.get(generated.len()).copied();
666 let after_opening = generated.get(opening_tokens.len()..).unwrap_or_default();
667 let delimiter_prefix_token_count =
668 delimiter_prefix_token_count(after_opening, delimiter_tokens);
669 let required_delimiter_token = opening_token
670 .or_else(|| delimiter_tokens.get(delimiter_prefix_token_count).copied())
671 .ok_or_else(|| {
672 FerrumError::internal("structured-output delimiter state has no next token")
673 })?;
674 let force_boundary = opening_token.is_some() || *forcing;
675 if force_boundary {
676 force_exact_token(logits, required_delimiter_token)?;
677 } else if let Some(hidden_control_token_ids) = hidden_control_token_ids {
678 for token_id in hidden_control_token_ids {
679 if required_delimiter_token == *token_id {
680 continue;
681 }
682 if let Some(logit) = logits.get_mut(*token_id as usize) {
683 *logit = f32::NEG_INFINITY;
684 }
685 }
686 }
687 return Ok(StructuredOutputMaskOutcome {
688 phase: if force_boundary {
689 StructuredOutputPhase::ForcingDelimiter
690 } else {
691 StructuredOutputPhase::WaitingForDelimiter
692 },
693 accepting: false,
694 liveness_intervention: false,
695 grammar_start_token_index: None,
696 required_delimiter_token_id: Some(required_delimiter_token),
697 grammar_owned_terminal_token_id: None,
698 });
699 }
700
701 let grammar_start_token_index = state.grammar_start.ok_or_else(|| {
702 FerrumError::internal(
703 "active structured-output processor has no grammar start token index",
704 )
705 })?;
706
707 let accepting = state.matcher.is_accepting().map_err(|error| {
708 FerrumError::model(format!(
709 "structured-output acceptance check failed: {error}"
710 ))
711 })?;
712 let mask = state.matcher.compute_mask_or_eos().map_err(|error| {
713 FerrumError::model(format!("structured-output mask failed: {error}"))
714 })?;
715 let mut finite_allowed = 0usize;
716 for (idx, logit) in logits.iter_mut().enumerate() {
717 let token = idx as u32;
718 let allowed_terminal = match &state.activation {
719 Activation::Harmony(boundary) => boundary.is_terminal(token),
720 _ => terminal_token_ids.is_some_and(|terminals| terminals.contains(&token)),
721 };
722 let hidden_non_terminal_control = hidden_control_token_ids
723 .is_some_and(|controls| controls.contains(&token))
724 && !allowed_terminal;
725 let protocol_control = matches!(&state.activation,
726 Activation::Harmony(boundary) if boundary.is_control(token) && !allowed_terminal);
727 let allowed = idx < self.vocab_size
728 && self.defined_token_ids.get(idx).copied().unwrap_or(false)
729 && !hidden_non_terminal_control
730 && !protocol_control
731 && (mask.is_allowed(token) || (accepting && allowed_terminal));
732 if !allowed {
733 *logit = f32::NEG_INFINITY;
734 } else if logit.is_finite() {
735 finite_allowed += 1;
736 }
737 }
738 if finite_allowed == 0 {
739 return Err(FerrumError::model(
740 "structured-output grammar has no legal finite token",
741 ));
742 }
743 let liveness_intervention = if !accepting
744 && state.trailing_identical_token_count >= self.liveness.max_identical_token_run
745 {
746 state
747 .trailing_grammar_token_id
748 .and_then(|token| logits.get_mut(token as usize))
749 .is_some_and(|logit| {
750 if logit.is_finite() && finite_allowed > 1 {
751 *logit = f32::NEG_INFINITY;
752 if state.last_liveness_intervention_at != Some(generated.len()) {
753 state.liveness_intervention_count += 1;
754 state.last_liveness_intervention_at = Some(generated.len());
755 }
756 true
757 } else {
758 false
759 }
760 })
761 } else {
762 false
763 };
764 Ok(StructuredOutputMaskOutcome {
765 phase: StructuredOutputPhase::EnforcingGrammar,
766 accepting,
767 liveness_intervention,
768 grammar_start_token_index: Some(grammar_start_token_index),
769 required_delimiter_token_id: None,
770 grammar_owned_terminal_token_id: None,
771 })
772 }
773
774 fn mask_undefined_token_ids(&self, logits: &mut [f32]) {
775 for (idx, logit) in logits.iter_mut().enumerate() {
776 if !self.defined_token_ids.get(idx).copied().unwrap_or(false) {
777 *logit = f32::NEG_INFINITY;
778 }
779 }
780 }
781
782 pub fn is_accepting(&self, generated: &[TokenId]) -> Result<bool> {
785 self.is_accepting_inner(generated, None)
786 }
787
788 pub fn is_accepting_with_terminals(
791 &self,
792 generated: &[TokenId],
793 terminal_token_ids: &HashSet<u32>,
794 ) -> Result<bool> {
795 self.is_accepting_inner(generated, Some(terminal_token_ids))
796 }
797
798 fn is_accepting_inner(
799 &self,
800 generated: &[TokenId],
801 terminal_token_ids: Option<&HashSet<u32>>,
802 ) -> Result<bool> {
803 Ok(self
804 .progress_inner(generated, terminal_token_ids)?
805 .accepting)
806 }
807
808 pub fn progress_with_terminals(
810 &self,
811 generated: &[TokenId],
812 terminal_token_ids: &HashSet<u32>,
813 ) -> Result<StructuredOutputProgress> {
814 self.progress_inner(generated, Some(terminal_token_ids))
815 }
816
817 fn progress_inner(
818 &self,
819 generated: &[TokenId],
820 terminal_token_ids: Option<&HashSet<u32>>,
821 ) -> Result<StructuredOutputProgress> {
822 let mut state = self.state.lock();
823 if state.composed.is_some() {
824 return auto_tools::progress(self, &mut state, generated, terminal_token_ids);
825 }
826 advance_state(&mut state, generated, terminal_token_ids, self.budget)?;
827 activate_forcing_if_due(&mut state, generated, self.budget);
828 let (phase, delimiter_token_count, delimiter_prefix_token_count, accepting) =
829 match &state.activation {
830 Activation::Boundary {
831 opening_tokens,
832 delimiter_tokens,
833 forcing,
834 ..
835 } => (
836 if generated.len() < opening_tokens.len() || *forcing {
837 StructuredOutputPhase::ForcingDelimiter
838 } else {
839 StructuredOutputPhase::WaitingForDelimiter
840 },
841 Some(opening_tokens.len() + delimiter_tokens.len()),
842 generated.len().min(opening_tokens.len())
843 + delimiter_prefix_token_count(
844 generated.get(opening_tokens.len()..).unwrap_or_default(),
845 delimiter_tokens,
846 ),
847 false,
848 ),
849 Activation::Active => (
850 StructuredOutputPhase::EnforcingGrammar,
851 None,
852 0,
853 state.matcher.is_accepting().map_err(|error| {
854 FerrumError::model(format!(
855 "structured-output acceptance check failed: {error}"
856 ))
857 })?,
858 ),
859 Activation::Harmony(boundary) => {
860 let (phase, length, prefix) = boundary.progress();
861 (
862 phase,
863 Some(length),
864 prefix,
865 boundary.in_payload()
866 && state.matcher.is_accepting().map_err(|error| {
867 FerrumError::model(format!(
868 "structured-output acceptance check failed: {error}"
869 ))
870 })?,
871 )
872 }
873 };
874 let grammar_tokens = state
875 .grammar_start
876 .and_then(|start| generated.get(start..))
877 .unwrap_or_default();
878 let trailing_token_id = grammar_tokens.last().map(|token| token.get());
879 let trailing_token_class = trailing_token_id.map(|token| {
880 self.json_token_classes
881 .get(token as usize)
882 .copied()
883 .unwrap_or(StructuredOutputTokenClass::Undefined)
884 });
885 let trailing_token_class_count = trailing_token_class.map_or(0, |class| {
886 grammar_tokens
887 .iter()
888 .rev()
889 .take_while(|token| {
890 self.json_token_classes
891 .get(token.get() as usize)
892 .copied()
893 .unwrap_or(StructuredOutputTokenClass::Undefined)
894 == class
895 })
896 .count()
897 });
898 let trailing_identical_token_count = trailing_token_id.map_or(0, |token_id| {
899 grammar_tokens
900 .iter()
901 .rev()
902 .take_while(|token| token.get() == token_id)
903 .count()
904 });
905 let opening_token_count = match &state.initial_activation {
906 Activation::Boundary { opening_tokens, .. } => opening_tokens.len(),
907 _ => 0,
908 };
909 Ok(StructuredOutputProgress {
910 phase,
911 generated_token_count: generated.len(),
912 consumed_token_count: state.consumed,
913 delimiter_token_count: delimiter_token_count
914 .or(self.budget.map(|budget| budget.boundary_token_count)),
915 delimiter_prefix_token_count,
916 reasoning_token_count: self.budget.map(|_| {
917 state
918 .boundary_start
919 .unwrap_or(generated.len())
920 .saturating_sub(opening_token_count)
921 }),
922 boundary_forced: state.boundary_forced,
923 budget: self.budget,
924 grammar_token_count: grammar_tokens.len(),
925 trailing_token_class,
926 trailing_token_class_count,
927 trailing_token_id,
928 trailing_identical_token_count,
929 liveness_identical_token_limit: self.liveness.max_identical_token_run,
930 liveness_intervention_count: state.liveness_intervention_count,
931 accepting,
932 })
933 }
934
935 pub fn reset(&self) -> Result<()> {
936 let mut state = self.state.lock();
937 if let Some(composed) = &state.composed {
938 state.matcher = composed.fresh_matcher();
942 } else {
943 state.matcher.reset().map_err(|error| {
944 FerrumError::internal(format!("reset structured output: {error}"))
945 })?;
946 }
947 state.activation = state.initial_activation.clone();
948 state.consumed = 0;
949 state.boundary_forced = false;
950 state.boundary_start = None;
951 state.grammar_start = matches!(state.initial_activation, Activation::Active).then_some(0);
952 state.trailing_grammar_token_id = None;
953 state.trailing_identical_token_count = 0;
954 state.liveness_intervention_count = 0;
955 state.last_liveness_intervention_at = None;
956 if let Some(composed) = &mut state.composed {
957 composed.reset();
958 state.grammar_start = composed.result_start();
959 }
960 Ok(())
961 }
962
963 pub fn is_protocol_complete_with_terminals(
966 &self,
967 generated: &[TokenId],
968 terminal_token_ids: &HashSet<u32>,
969 ) -> Result<bool> {
970 let mut state = self.state.lock();
971 if !state
972 .composed
973 .as_ref()
974 .is_some_and(|composed| composed.is_native())
975 {
976 return Ok(false);
977 }
978 auto_tools::advance(&mut state, generated, Some(terminal_token_ids))?;
979 state.matcher.is_accepting().map_err(|error| {
980 FerrumError::model(format!(
981 "structured-output acceptance check failed: {error}"
982 ))
983 })
984 }
985
986 pub fn complete_root_text_len_with_terminals(
990 &self,
991 generated: &[TokenId],
992 terminal_token_ids: &HashSet<u32>,
993 ) -> Result<Option<usize>> {
994 let mut state = self.state.lock();
995 if state.composed.is_none() {
996 return Ok(None);
997 }
998 auto_tools::advance(&mut state, generated, Some(terminal_token_ids))?;
999 let accepting = state.matcher.is_accepting().map_err(|error| {
1000 FerrumError::model(format!(
1001 "structured-output acceptance check failed: {error}"
1002 ))
1003 })?;
1004 Ok(accepting.then(|| {
1005 state
1006 .composed
1007 .as_ref()
1008 .expect("composed processor")
1009 .complete_text_len()
1010 }))
1011 }
1012
1013 pub fn classified_result_with_terminals(
1017 &self,
1018 generated: &[TokenId],
1019 terminal_token_ids: &HashSet<u32>,
1020 ) -> Result<Option<(ferrum_types::StructuredOutputBranch, String)>> {
1021 let mut state = self.state.lock();
1022 if state.composed.is_none() {
1023 return Ok(None);
1024 }
1025 auto_tools::classify(&mut state, generated, terminal_token_ids)
1026 }
1027}
1028
1029fn activate_forcing_if_due(
1030 state: &mut ProcessorState,
1031 generated: &[TokenId],
1032 budget: Option<StructuredOutputBudgetPlan>,
1033) {
1034 if let Activation::Harmony(boundary) = &mut state.activation {
1035 boundary.activate_forcing_if_due(generated.len());
1036 state.boundary_forced = boundary.is_forced();
1037 return;
1038 }
1039 let Some(budget) = budget else {
1040 return;
1041 };
1042 let should_force = match &state.activation {
1043 Activation::Boundary {
1044 opening_tokens,
1045 allow_reasoning,
1046 forcing,
1047 ..
1048 } => {
1049 !forcing
1050 && generated.len() >= opening_tokens.len()
1051 && (!allow_reasoning
1052 || generated.len() - opening_tokens.len() >= budget.reasoning_token_limit)
1053 }
1054 _ => false,
1055 };
1056 if should_force {
1057 let delimiter_prefix_token_count = match &state.activation {
1058 Activation::Boundary {
1059 opening_tokens,
1060 delimiter_tokens,
1061 ..
1062 } => delimiter_prefix_token_count(&generated[opening_tokens.len()..], delimiter_tokens),
1063 Activation::Active | Activation::Harmony(_) => 0,
1064 };
1065 if let Activation::Boundary { forcing, .. } = &mut state.activation {
1066 *forcing = true;
1067 }
1068 state.boundary_forced = true;
1069 state.boundary_start = Some(generated.len() - delimiter_prefix_token_count);
1070 }
1071}
1072
1073fn force_exact_token(logits: &mut [f32], required_token: u32) -> Result<()> {
1074 let required_index = required_token as usize;
1075 if required_index >= logits.len() {
1076 return Err(FerrumError::model(format!(
1077 "structured-output delimiter token {required_token} is outside logits width {}",
1078 logits.len()
1079 )));
1080 }
1081 logits.fill(f32::NEG_INFINITY);
1082 logits[required_index] = 0.0;
1083 Ok(())
1084}
1085
1086fn delimiter_prefix_token_count(generated: &[TokenId], delimiter_tokens: &[u32]) -> usize {
1087 let max_prefix = generated
1088 .len()
1089 .min(delimiter_tokens.len().saturating_sub(1));
1090 (1..=max_prefix)
1091 .rev()
1092 .find(|prefix_len| {
1093 generated[generated.len() - prefix_len..]
1094 .iter()
1095 .zip(&delimiter_tokens[..*prefix_len])
1096 .all(|(token, expected)| token.get() == *expected)
1097 })
1098 .unwrap_or(0)
1099}
1100
1101fn advance_state(
1102 state: &mut ProcessorState,
1103 generated: &[TokenId],
1104 terminal_token_ids: Option<&HashSet<u32>>,
1105 budget: Option<StructuredOutputBudgetPlan>,
1106) -> Result<()> {
1107 if state.consumed > generated.len() {
1108 return Err(FerrumError::internal(
1109 "structured-output token history moved backwards without reset",
1110 ));
1111 }
1112
1113 if let Activation::Harmony(boundary) = &mut state.activation {
1114 while state.consumed < generated.len() && !boundary.in_payload() {
1115 let index = state.consumed;
1116 if let Some(start) = boundary.observe(generated[index].get(), index)? {
1117 state.grammar_start = Some(start);
1118 state.boundary_start = Some(boundary.boundary_start());
1119 }
1120 state.consumed = index + 1;
1121 }
1122 state.boundary_forced = boundary.is_forced();
1123 if !boundary.in_payload() {
1124 return Ok(());
1125 }
1126 }
1127
1128 if let Activation::Boundary {
1129 opening_tokens,
1130 delimiter_tokens,
1131 allow_reasoning,
1132 ..
1133 } = &state.activation
1134 {
1135 if opening_tokens
1136 .iter()
1137 .zip(generated)
1138 .any(|(expected, token)| *expected != token.get())
1139 {
1140 return Err(FerrumError::model(
1141 "structured-output reasoning envelope has an invalid opening header",
1142 ));
1143 }
1144 if generated.len() < opening_tokens.len() {
1145 state.consumed = generated.len();
1146 return Ok(());
1147 }
1148 if !allow_reasoning
1149 && delimiter_tokens
1150 .iter()
1151 .zip(&generated[opening_tokens.len()..])
1152 .any(|(expected, token)| *expected != token.get())
1153 {
1154 return Err(FerrumError::model(
1155 "structured-output reasoning is disabled; the envelope must close immediately",
1156 ));
1157 }
1158 let search_from = state
1159 .consumed
1160 .saturating_sub(delimiter_tokens.len())
1161 .max(opening_tokens.len());
1162 if let Some(offset) = generated[search_from..]
1163 .windows(delimiter_tokens.len())
1164 .position(|window| {
1165 window
1166 .iter()
1167 .zip(delimiter_tokens)
1168 .all(|(token, expected)| token.get() == *expected)
1169 })
1170 {
1171 let grammar_start = search_from + offset + delimiter_tokens.len();
1172 if !opening_tokens.is_empty() {
1173 let reasoning_count = grammar_start - delimiter_tokens.len() - opening_tokens.len();
1174 state.boundary_forced |= !allow_reasoning
1175 || budget.is_some_and(|plan| {
1176 reasoning_count + delimiter_tokens.len() - 1 >= plan.reasoning_token_limit
1177 });
1178 }
1179 state.boundary_start = Some(grammar_start - delimiter_tokens.len());
1180 state.grammar_start = Some(grammar_start);
1181 state.activation = Activation::Active;
1182 state.consumed = grammar_start;
1183 state.trailing_grammar_token_id = None;
1184 state.trailing_identical_token_count = 0;
1185 } else {
1186 state.consumed = generated.len();
1187 return Ok(());
1188 }
1189 }
1190
1191 for token in &generated[state.consumed..] {
1192 if let Activation::Harmony(boundary) = &mut state.activation {
1193 let accepting = state.matcher.is_accepting().map_err(|error| {
1194 FerrumError::model(format!(
1195 "structured-output acceptance check failed: {error}"
1196 ))
1197 })?;
1198 if boundary.observe_payload_control(token.get(), accepting)? {
1199 continue;
1200 }
1201 } else if terminal_token_ids.is_some_and(|terminals| terminals.contains(&token.get()))
1202 && state.matcher.is_accepting().map_err(|error| {
1203 FerrumError::model(format!(
1204 "structured-output acceptance check failed: {error}"
1205 ))
1206 })?
1207 {
1208 continue;
1209 }
1210 state.matcher.consume_token(token.get()).map_err(|error| {
1211 FerrumError::model(format!(
1212 "structured-output token {} violated the grammar: {error}",
1213 token.get()
1214 ))
1215 })?;
1216 if state.trailing_grammar_token_id == Some(token.get()) {
1217 state.trailing_identical_token_count += 1;
1218 } else {
1219 state.trailing_grammar_token_id = Some(token.get());
1220 state.trailing_identical_token_count = 1;
1221 }
1222 }
1223 state.consumed = generated.len();
1224 Ok(())
1225}
1226
1227struct FerrumTokenizerEnv {
1228 tokenizer: Arc<dyn Tokenizer + Send + Sync>,
1229 trie: TokTrie,
1230}
1231
1232impl TokenizerEnv for FerrumTokenizerEnv {
1233 fn tok_trie(&self) -> &TokTrie {
1234 &self.trie
1235 }
1236
1237 fn tokenize_bytes(&self, bytes: &[u8]) -> Vec<u32> {
1238 str::from_utf8(bytes)
1239 .ok()
1240 .and_then(|text| self.tokenizer.encode(text, false).ok())
1241 .map(|tokens| tokens.into_iter().map(|token| token.get()).collect())
1242 .unwrap_or_else(|| self.trie.greedy_tokenize(bytes))
1243 }
1244
1245 fn tokenize_is_canonical(&self) -> bool {
1246 false
1247 }
1248}
1249
1250fn tokenizer_special_ids(tokenizer: &(dyn Tokenizer + Send + Sync)) -> HashSet<u32> {
1251 let special = tokenizer.special_tokens();
1252 [
1253 special.bos_token,
1254 special.eos_token,
1255 special.unk_token,
1256 special.pad_token,
1257 special.sep_token,
1258 special.cls_token,
1259 special.mask_token,
1260 ]
1261 .into_iter()
1262 .flatten()
1263 .chain(special.extra_eos_tokens.iter().copied())
1264 .map(|token| token.get())
1265 .collect()
1266}
1267
1268fn special_token_marker(token: TokenId) -> Vec<u8> {
1269 let mut marker = vec![TokTrie::SPECIAL_TOKEN_MARKER];
1270 marker.extend_from_slice(format!("[{}]", token.get()).as_bytes());
1271 marker
1272}
1273
1274fn classify_json_token_bytes(bytes: &[u8]) -> StructuredOutputTokenClass {
1275 if bytes.is_empty() {
1276 StructuredOutputTokenClass::Undefined
1277 } else if bytes
1278 .iter()
1279 .all(|byte| matches!(byte, b' ' | b'\n' | b'\r' | b'\t'))
1280 {
1281 StructuredOutputTokenClass::Whitespace
1282 } else if bytes
1283 .iter()
1284 .all(|byte| byte.is_ascii_digit() || matches!(byte, b'-' | b'+' | b'.' | b'e' | b'E'))
1285 {
1286 StructuredOutputTokenClass::Number
1287 } else if bytes
1288 .iter()
1289 .all(|byte| matches!(byte, b'{' | b'}' | b'[' | b']' | b',' | b':'))
1290 {
1291 StructuredOutputTokenClass::Structural
1292 } else if bytes.iter().all(|byte| matches!(byte, b'"' | b'\\')) {
1293 StructuredOutputTokenClass::StringBoundary
1294 } else if bytes.iter().all(u8::is_ascii_alphabetic) {
1295 StructuredOutputTokenClass::Literal
1296 } else {
1297 StructuredOutputTokenClass::Other
1298 }
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303 use super::*;
1304 use ferrum_interfaces::tokenizer::{ChatMessage, TokenizerInfo, TokenizerType};
1305 use ferrum_types::SpecialTokens;
1306
1307 const EOS: u32 = 256;
1308 const TEST_MAX_OUTPUT_TOKENS: usize = 128;
1309
1310 pub(super) struct ByteTokenizer {
1311 special: SpecialTokens,
1312 token_text: Vec<String>,
1313 }
1314
1315 impl ByteTokenizer {
1316 pub(super) fn new() -> Self {
1317 let mut token_text = (0u16..=255)
1318 .map(|byte| char::from_u32(byte as u32).unwrap().to_string())
1319 .collect::<Vec<_>>();
1320 token_text.push("<eos>".to_string());
1321 Self {
1322 special: SpecialTokens {
1323 eos_token: Some(TokenId::new(EOS)),
1324 ..SpecialTokens::default()
1325 },
1326 token_text,
1327 }
1328 }
1329 }
1330
1331 impl Tokenizer for ByteTokenizer {
1332 fn encode(&self, text: &str, _add_special: bool) -> Result<Vec<TokenId>> {
1333 Ok(text
1334 .as_bytes()
1335 .iter()
1336 .map(|byte| TokenId::new(*byte as u32))
1337 .collect())
1338 }
1339
1340 fn decode(&self, tokens: &[TokenId], _skip_special: bool) -> Result<String> {
1341 Ok(tokens
1342 .iter()
1343 .filter(|token| token.get() < 256)
1344 .map(|token| token.get() as u8 as char)
1345 .collect())
1346 }
1347
1348 fn decode_incremental(&self, _prev: &[TokenId], next: TokenId) -> Result<String> {
1349 self.decode(&[next], true)
1350 }
1351
1352 fn vocab_size(&self) -> usize {
1353 self.token_text.len()
1354 }
1355
1356 fn special_tokens(&self) -> &SpecialTokens {
1357 &self.special
1358 }
1359
1360 fn token_id(&self, text: &str) -> Option<TokenId> {
1361 (text.len() == 1).then(|| TokenId::new(text.as_bytes()[0] as u32))
1362 }
1363
1364 fn token_text(&self, token_id: TokenId) -> Option<&str> {
1365 self.token_text
1366 .get(token_id.get() as usize)
1367 .map(String::as_str)
1368 }
1369
1370 fn apply_chat_template(&self, messages: &[ChatMessage]) -> Result<String> {
1371 Ok(messages
1372 .iter()
1373 .map(|message| message.content.as_str())
1374 .collect::<Vec<_>>()
1375 .join("\n"))
1376 }
1377
1378 fn info(&self) -> TokenizerInfo {
1379 TokenizerInfo {
1380 tokenizer_type: TokenizerType::BPE,
1381 vocab_size: self.vocab_size(),
1382 special_tokens: self.special.clone(),
1383 supports_incremental: true,
1384 supports_chat_template: false,
1385 max_token_length: Some(1),
1386 model_name: Some("byte-test".to_string()),
1387 }
1388 }
1389 }
1390
1391 struct MergedObjectTokenizer {
1392 inner: ByteTokenizer,
1393 }
1394
1395 impl MergedObjectTokenizer {
1396 const OBJECT: u32 = 256;
1397 const EOS: u32 = 257;
1398
1399 fn new() -> Self {
1400 let mut inner = ByteTokenizer::new();
1401 inner.token_text[EOS as usize] = "{}".to_string();
1402 inner.token_text.push("<eos>".to_string());
1403 inner.special.eos_token = Some(TokenId::new(Self::EOS));
1404 Self { inner }
1405 }
1406 }
1407
1408 impl Tokenizer for MergedObjectTokenizer {
1409 fn encode(&self, text: &str, add_special: bool) -> Result<Vec<TokenId>> {
1410 self.inner.encode(text, add_special)
1411 }
1412
1413 fn decode(&self, tokens: &[TokenId], skip_special: bool) -> Result<String> {
1414 let mut decoded = String::new();
1415 for token in tokens {
1416 match token.get() {
1417 Self::OBJECT => decoded.push_str("{}"),
1418 Self::EOS if skip_special => {}
1419 Self::EOS => decoded.push_str("<eos>"),
1420 _ => decoded.push_str(&self.inner.decode(&[*token], skip_special)?),
1421 }
1422 }
1423 Ok(decoded)
1424 }
1425
1426 fn decode_incremental(&self, _prev: &[TokenId], next: TokenId) -> Result<String> {
1427 self.decode(&[next], true)
1428 }
1429
1430 fn vocab_size(&self) -> usize {
1431 self.inner.token_text.len()
1432 }
1433
1434 fn special_tokens(&self) -> &SpecialTokens {
1435 &self.inner.special
1436 }
1437
1438 fn token_id(&self, text: &str) -> Option<TokenId> {
1439 (text == "{}")
1440 .then(|| TokenId::new(Self::OBJECT))
1441 .or_else(|| self.inner.token_id(text))
1442 }
1443
1444 fn token_text(&self, token_id: TokenId) -> Option<&str> {
1445 self.inner.token_text(token_id)
1446 }
1447
1448 fn info(&self) -> TokenizerInfo {
1449 TokenizerInfo {
1450 vocab_size: self.vocab_size(),
1451 max_token_length: Some(2),
1452 model_name: Some("merged-object-test".to_string()),
1453 ..self.inner.info()
1454 }
1455 }
1456 }
1457
1458 fn factory() -> StructuredOutputFactory {
1459 StructuredOutputFactory::new(Arc::new(ByteTokenizer::new())).unwrap()
1460 }
1461
1462 fn assert_and_append(
1463 processor: &StructuredOutputProcessor,
1464 generated: &mut Vec<TokenId>,
1465 text: &str,
1466 ) {
1467 for byte in text.bytes() {
1468 let mut logits = vec![0.0; EOS as usize + 1];
1469 processor.mask_logits(&mut logits, generated).unwrap();
1470 assert!(
1471 logits[byte as usize].is_finite(),
1472 "byte {byte:?} rejected after {:?}",
1473 generated
1474 );
1475 generated.push(TokenId::new(byte as u32));
1476 }
1477 }
1478
1479 #[test]
1480 fn json_object_hard_masks_non_object_roots() {
1481 let processor = factory()
1482 .create_processor(
1483 &ResponseFormat::JsonObject,
1484 &StructuredOutputStart::Immediate,
1485 TEST_MAX_OUTPUT_TOKENS,
1486 &HashSet::new(),
1487 &[],
1488 )
1489 .unwrap()
1490 .unwrap();
1491 let mut logits = vec![0.0; EOS as usize + 1];
1492 processor.mask_logits(&mut logits, &[]).unwrap();
1493 assert!(logits[b'{' as usize].is_finite());
1494 assert!(!logits[b'[' as usize].is_finite());
1495 assert!(!logits[b'`' as usize].is_finite());
1496 assert!(!logits[EOS as usize].is_finite());
1497 }
1498
1499 #[test]
1500 fn json_object_uses_compact_separators_without_unbounded_whitespace() {
1501 let processor = factory()
1502 .create_processor(
1503 &ResponseFormat::JsonObject,
1504 &StructuredOutputStart::Immediate,
1505 TEST_MAX_OUTPUT_TOKENS,
1506 &HashSet::new(),
1507 &[],
1508 )
1509 .unwrap()
1510 .unwrap();
1511 let generated = vec![TokenId::new(b'{' as u32)];
1512 let mut logits = vec![0.0; EOS as usize + 1];
1513 processor.mask_logits(&mut logits, &generated).unwrap();
1514
1515 assert!(!logits[b' ' as usize].is_finite());
1516 assert!(logits[b'}' as usize].is_finite());
1517 assert!(logits[b'"' as usize].is_finite());
1518 }
1519
1520 #[test]
1521 fn undefined_model_vocab_ids_are_masked_inside_wildcard_strings() {
1522 let tokenizer = Arc::new(ByteTokenizer::new());
1523 let undefined_token = tokenizer.vocab_size() as u32;
1524 assert_eq!(
1525 tokenizer.token_bytes(TokenId::new(undefined_token)),
1526 Some(Vec::new()),
1527 "the test tokenizer must reproduce a decoder that returns empty text for an unknown id"
1528 );
1529 let processor = StructuredOutputFactory::new_with_model_vocab_size(
1530 tokenizer,
1531 Some(undefined_token as usize + 1),
1532 )
1533 .unwrap()
1534 .create_processor(
1535 &ResponseFormat::JsonObject,
1536 &StructuredOutputStart::Immediate,
1537 TEST_MAX_OUTPUT_TOKENS,
1538 &HashSet::new(),
1539 &[],
1540 )
1541 .unwrap()
1542 .unwrap();
1543
1544 let mut generated = Vec::new();
1545 assert_and_append(&processor, &mut generated, r#"{"value":"Ferrum "#);
1546 let mut logits = vec![0.0; undefined_token as usize + 1];
1547 processor.mask_logits(&mut logits, &generated).unwrap();
1548 assert!(logits[b'x' as usize].is_finite());
1549 assert!(!logits[undefined_token as usize].is_finite());
1550 }
1551
1552 #[test]
1553 fn json_object_accepts_nested_unicode_escape_and_eos_only_after_close() {
1554 let processor = factory()
1555 .create_processor(
1556 &ResponseFormat::JsonObject,
1557 &StructuredOutputStart::Immediate,
1558 TEST_MAX_OUTPUT_TOKENS,
1559 &HashSet::new(),
1560 &[],
1561 )
1562 .unwrap()
1563 .unwrap();
1564 let mut generated = Vec::new();
1565 assert_and_append(
1566 &processor,
1567 &mut generated,
1568 r#"{"items":[true,null,{"name":"line\u000A"}],"n":-1.2e+3}"#,
1569 );
1570 assert!(processor.is_accepting(&generated).unwrap());
1571 let mut logits = vec![0.0; EOS as usize + 1];
1572 processor.mask_logits(&mut logits, &generated).unwrap();
1573 assert!(logits[EOS as usize].is_finite());
1574 assert!(!logits[b'x' as usize].is_finite());
1575 }
1576
1577 #[test]
1578 fn json_object_accepts_a_complete_root_from_one_merged_token() {
1579 let processor = StructuredOutputFactory::new(Arc::new(MergedObjectTokenizer::new()))
1580 .unwrap()
1581 .create_processor(
1582 &ResponseFormat::JsonObject,
1583 &StructuredOutputStart::Immediate,
1584 TEST_MAX_OUTPUT_TOKENS,
1585 &HashSet::new(),
1586 &[],
1587 )
1588 .unwrap()
1589 .unwrap();
1590 let generated = vec![TokenId::new(MergedObjectTokenizer::OBJECT)];
1591 let terminals = HashSet::from([MergedObjectTokenizer::EOS]);
1592
1593 let progress = processor
1594 .progress_with_terminals(&generated, &terminals)
1595 .unwrap();
1596 assert!(progress.accepting);
1597
1598 let mut logits = vec![0.0; MergedObjectTokenizer::EOS as usize + 1];
1599 let outcome = processor
1600 .mask_logits_with_terminals(&mut logits, &generated, &terminals, &HashSet::new())
1601 .unwrap();
1602 assert!(outcome.accepting);
1603 assert_eq!(outcome.grammar_start_token_index, Some(0));
1604 assert!(!logits[MergedObjectTokenizer::OBJECT as usize].is_finite());
1605 assert!(logits[MergedObjectTokenizer::EOS as usize].is_finite());
1606 }
1607
1608 #[test]
1609 fn json_object_breaks_an_unbounded_identical_token_run_when_closure_is_legal() {
1610 let processor = StructuredOutputFactory::new(Arc::new(MergedObjectTokenizer::new()))
1611 .unwrap()
1612 .create_processor(
1613 &ResponseFormat::JsonObject,
1614 &StructuredOutputStart::Immediate,
1615 64,
1616 &HashSet::new(),
1617 &[],
1618 )
1619 .unwrap()
1620 .unwrap();
1621 let mut generated =
1622 br#"{"marker":""#.iter().map(|byte| TokenId::new(*byte as u32)).collect::<Vec<_>>();
1623 generated.extend(std::iter::repeat_n(
1624 TokenId::new(MergedObjectTokenizer::OBJECT),
1625 32,
1626 ));
1627
1628 let mut logits = vec![0.0; MergedObjectTokenizer::EOS as usize + 1];
1629 let outcome = processor
1630 .mask_logits_with_terminals(
1631 &mut logits,
1632 &generated,
1633 &HashSet::from([MergedObjectTokenizer::EOS]),
1634 &HashSet::new(),
1635 )
1636 .unwrap();
1637
1638 assert!(!outcome.accepting);
1639 assert!(outcome.liveness_intervention);
1640 assert!(!logits[MergedObjectTokenizer::OBJECT as usize].is_finite());
1641 assert!(logits[b'"' as usize].is_finite());
1642
1643 generated.extend([TokenId::new(b'"' as u32), TokenId::new(b'}' as u32)]);
1644 let progress = processor
1645 .progress_with_terminals(&generated, &HashSet::from([MergedObjectTokenizer::EOS]))
1646 .unwrap();
1647 assert!(progress.accepting);
1648 assert_eq!(progress.liveness_identical_token_limit, 32);
1649 assert_eq!(progress.liveness_intervention_count, 1);
1650 }
1651
1652 #[test]
1653 fn structured_liveness_guard_preserves_the_only_finite_grammar_candidate() {
1654 let processor = StructuredOutputFactory::new(Arc::new(MergedObjectTokenizer::new()))
1655 .unwrap()
1656 .create_processor(
1657 &ResponseFormat::JsonObject,
1658 &StructuredOutputStart::Immediate,
1659 64,
1660 &HashSet::new(),
1661 &[],
1662 )
1663 .unwrap()
1664 .unwrap();
1665 let mut generated =
1666 br#"{"marker":""#.iter().map(|byte| TokenId::new(*byte as u32)).collect::<Vec<_>>();
1667 generated.extend(std::iter::repeat_n(
1668 TokenId::new(MergedObjectTokenizer::OBJECT),
1669 32,
1670 ));
1671 let mut logits = vec![f32::NEG_INFINITY; MergedObjectTokenizer::EOS as usize + 1];
1672 logits[MergedObjectTokenizer::OBJECT as usize] = 0.0;
1673
1674 let outcome = processor
1675 .mask_logits_with_terminals(
1676 &mut logits,
1677 &generated,
1678 &HashSet::from([MergedObjectTokenizer::EOS]),
1679 &HashSet::new(),
1680 )
1681 .unwrap();
1682
1683 assert!(!outcome.liveness_intervention);
1684 assert!(logits[MergedObjectTokenizer::OBJECT as usize].is_finite());
1685 }
1686
1687 #[test]
1688 fn structured_liveness_limit_is_derived_from_the_guaranteed_result_budget() {
1689 let budget = StructuredOutputBudgetPlan::automatic(4096, 1).unwrap();
1690 assert_eq!(budget.structured_reserve_tokens, 1024);
1691 assert_eq!(
1692 StructuredOutputLivenessPolicy::for_request(4096, Some(budget)).max_identical_token_run,
1693 512
1694 );
1695 assert_eq!(
1696 StructuredOutputLivenessPolicy::for_request(64, None).max_identical_token_run,
1697 32
1698 );
1699 }
1700
1701 #[test]
1702 fn strict_schema_rejects_wrong_property_and_accepts_required_value() {
1703 let schema = r#"{
1704 "type":"object",
1705 "properties":{"answer":{"const":42}},
1706 "required":["answer"],
1707 "additionalProperties":false
1708 }"#;
1709 let processor = factory()
1710 .create_processor(
1711 &ResponseFormat::JsonSchema(schema.to_string()),
1712 &StructuredOutputStart::Immediate,
1713 TEST_MAX_OUTPUT_TOKENS,
1714 &HashSet::new(),
1715 &[],
1716 )
1717 .unwrap()
1718 .unwrap();
1719 let mut generated = Vec::new();
1720 assert_and_append(&processor, &mut generated, r#"{"answer":42}"#);
1721 assert!(processor.is_accepting(&generated).unwrap());
1722 }
1723
1724 #[test]
1725 fn request_cannot_override_compact_json_compiler_policy() {
1726 let schema = r#"{
1727 "type":"object",
1728 "properties":{"answer":{"const":42}},
1729 "required":["answer"],
1730 "additionalProperties":false,
1731 "x-guidance":{
1732 "item_separator":", ",
1733 "key_separator":": ",
1734 "whitespace_flexible":true
1735 }
1736 }"#;
1737 let processor = factory()
1738 .create_processor(
1739 &ResponseFormat::JsonSchema(schema.to_string()),
1740 &StructuredOutputStart::Immediate,
1741 TEST_MAX_OUTPUT_TOKENS,
1742 &HashSet::new(),
1743 &[],
1744 )
1745 .unwrap()
1746 .unwrap();
1747 let mut generated = Vec::new();
1748 assert_and_append(&processor, &mut generated, r#"{"answer":"#);
1749 let mut logits = vec![0.0; EOS as usize + 1];
1750 processor.mask_logits(&mut logits, &generated).unwrap();
1751
1752 assert!(!logits[b' ' as usize].is_finite());
1753 assert!(logits[b'4' as usize].is_finite());
1754 assert_and_append(&processor, &mut generated, "42}");
1755 assert!(processor.is_accepting(&generated).unwrap());
1756 }
1757
1758 #[test]
1759 fn boolean_json_schema_keeps_its_semantics_under_compact_policy() {
1760 let processor = factory()
1761 .create_processor(
1762 &ResponseFormat::JsonSchema("true".to_string()),
1763 &StructuredOutputStart::Immediate,
1764 TEST_MAX_OUTPUT_TOKENS,
1765 &HashSet::new(),
1766 &[],
1767 )
1768 .unwrap()
1769 .unwrap();
1770 let mut generated = Vec::new();
1771 assert_and_append(&processor, &mut generated, "true");
1772 assert!(processor.is_accepting(&generated).unwrap());
1773 }
1774
1775 #[test]
1776 fn terminal_progress_classifies_an_unclosed_number_without_retaining_text() {
1777 let processor = factory()
1778 .create_processor(
1779 &ResponseFormat::JsonSchema(
1780 r#"{"type":"object","properties":{"value":{"type":"integer"}},"required":["value"],"additionalProperties":false}"#
1781 .to_string(),
1782 ),
1783 &StructuredOutputStart::Immediate,
1784 TEST_MAX_OUTPUT_TOKENS,
1785 &HashSet::new(),
1786 &[],
1787 )
1788 .unwrap()
1789 .unwrap();
1790 let generated =
1791 r#"{"value":123777"#.bytes().map(|byte| TokenId::new(byte as u32)).collect::<Vec<_>>();
1792 let progress = processor
1793 .progress_with_terminals(&generated, &HashSet::new())
1794 .unwrap();
1795
1796 assert_eq!(progress.phase, StructuredOutputPhase::EnforcingGrammar);
1797 assert_eq!(progress.grammar_token_count, generated.len());
1798 assert_eq!(
1799 progress.trailing_token_class,
1800 Some(StructuredOutputTokenClass::Number)
1801 );
1802 assert_eq!(progress.trailing_token_class_count, 6);
1803 assert_eq!(progress.trailing_token_id, Some(b'7' as u32));
1804 assert_eq!(progress.trailing_identical_token_count, 3);
1805 assert!(!progress.accepting);
1806 }
1807
1808 #[test]
1809 fn lexical_diagnostics_classify_json_bytes_without_decoding_content() {
1810 assert_eq!(
1811 classify_json_token_bytes(b" \n\t"),
1812 StructuredOutputTokenClass::Whitespace
1813 );
1814 assert_eq!(
1815 classify_json_token_bytes(b"-12.5e+3"),
1816 StructuredOutputTokenClass::Number
1817 );
1818 assert_eq!(
1819 classify_json_token_bytes(br#"{}[],:"#),
1820 StructuredOutputTokenClass::Structural
1821 );
1822 assert_eq!(
1823 classify_json_token_bytes(br#"\""#),
1824 StructuredOutputTokenClass::StringBoundary
1825 );
1826 assert_eq!(
1827 classify_json_token_bytes(b"true"),
1828 StructuredOutputTokenClass::Literal
1829 );
1830 assert_eq!(
1831 classify_json_token_bytes(br#""value":"#),
1832 StructuredOutputTokenClass::Other
1833 );
1834 }
1835
1836 struct FragmentedUtf8Tokenizer {
1837 special: SpecialTokens,
1838 token_text: Vec<String>,
1839 }
1840
1841 impl FragmentedUtf8Tokenizer {
1842 const FIRE_HEAD: u32 = 128;
1843 const FIRE_TAIL: u32 = 129;
1844 const EOS: u32 = 130;
1845
1846 fn new() -> Self {
1847 let mut token_text = (0u8..=127)
1848 .map(|byte| (byte as char).to_string())
1849 .collect::<Vec<_>>();
1850 token_text.extend(["\u{fffd}".to_string(), "\u{fffd}".to_string()]);
1851 token_text.push("<eos>".to_string());
1852 Self {
1853 special: SpecialTokens {
1854 eos_token: Some(TokenId::new(Self::EOS)),
1855 ..SpecialTokens::default()
1856 },
1857 token_text,
1858 }
1859 }
1860
1861 fn raw_bytes(token: TokenId) -> Option<Vec<u8>> {
1862 match token.get() {
1863 byte @ 0..=127 => Some(vec![byte as u8]),
1864 Self::FIRE_HEAD => Some(vec![0xf0, 0x9f]),
1865 Self::FIRE_TAIL => Some(vec![0x94, 0xa5]),
1866 _ => None,
1867 }
1868 }
1869 }
1870
1871 impl Tokenizer for FragmentedUtf8Tokenizer {
1872 fn encode(&self, text: &str, _add_special: bool) -> Result<Vec<TokenId>> {
1873 let mut tokens = Vec::new();
1874 let mut bytes = text.as_bytes();
1875 while let Some((&byte, remaining)) = bytes.split_first() {
1876 if bytes.starts_with(&[0xf0, 0x9f, 0x94, 0xa5]) {
1877 tokens.push(TokenId::new(Self::FIRE_HEAD));
1878 tokens.push(TokenId::new(Self::FIRE_TAIL));
1879 bytes = &bytes[4..];
1880 } else if byte <= 127 {
1881 tokens.push(TokenId::new(byte as u32));
1882 bytes = remaining;
1883 } else {
1884 return Err(FerrumError::tokenizer(
1885 "fragmented UTF-8 test tokenizer received unsupported input",
1886 ));
1887 }
1888 }
1889 Ok(tokens)
1890 }
1891
1892 fn decode(&self, tokens: &[TokenId], _skip_special: bool) -> Result<String> {
1893 let bytes = tokens
1894 .iter()
1895 .filter_map(|token| Self::raw_bytes(*token))
1896 .flatten()
1897 .collect::<Vec<_>>();
1898 Ok(String::from_utf8_lossy(&bytes).into_owned())
1899 }
1900
1901 fn decode_incremental(&self, _prev: &[TokenId], next: TokenId) -> Result<String> {
1902 self.decode(&[next], true)
1903 }
1904
1905 fn vocab_size(&self) -> usize {
1906 self.token_text.len()
1907 }
1908
1909 fn special_tokens(&self) -> &SpecialTokens {
1910 &self.special
1911 }
1912
1913 fn token_id(&self, text: &str) -> Option<TokenId> {
1914 (text.len() == 1 && text.is_ascii()).then(|| TokenId::new(text.as_bytes()[0] as u32))
1915 }
1916
1917 fn token_text(&self, token_id: TokenId) -> Option<&str> {
1918 self.token_text
1919 .get(token_id.get() as usize)
1920 .map(String::as_str)
1921 }
1922
1923 fn token_bytes(&self, token_id: TokenId) -> Option<Vec<u8>> {
1924 Self::raw_bytes(token_id)
1925 }
1926
1927 fn info(&self) -> TokenizerInfo {
1928 TokenizerInfo {
1929 tokenizer_type: TokenizerType::BPE,
1930 vocab_size: self.vocab_size(),
1931 special_tokens: self.special.clone(),
1932 supports_incremental: true,
1933 supports_chat_template: false,
1934 max_token_length: Some(2),
1935 model_name: Some("fragmented-utf8-test".to_string()),
1936 }
1937 }
1938 }
1939
1940 #[test]
1941 fn strict_schema_accepts_utf8_split_across_byte_level_tokens() {
1942 let tokenizer = Arc::new(FragmentedUtf8Tokenizer::new());
1943 assert!(tokenizer
1944 .decode(&[TokenId::new(FragmentedUtf8Tokenizer::FIRE_HEAD)], false)
1945 .unwrap()
1946 .contains('\u{fffd}'));
1947 let processor = StructuredOutputFactory::new(tokenizer)
1948 .unwrap()
1949 .create_processor(
1950 &ResponseFormat::JsonSchema(
1951 r#"{"type":"object","properties":{"value":{"const":"\ud83d\udd25"}},"required":["value"],"additionalProperties":false}"#
1952 .to_string(),
1953 ),
1954 &StructuredOutputStart::Immediate,
1955 TEST_MAX_OUTPUT_TOKENS,
1956 &HashSet::new(),
1957 &[],
1958 )
1959 .unwrap()
1960 .unwrap();
1961
1962 let mut generated = Vec::new();
1963 for token in r#"{"value":""#
1964 .bytes()
1965 .map(|byte| TokenId::new(byte as u32))
1966 .chain([
1967 TokenId::new(FragmentedUtf8Tokenizer::FIRE_HEAD),
1968 TokenId::new(FragmentedUtf8Tokenizer::FIRE_TAIL),
1969 ])
1970 .chain(r#""}"#.bytes().map(|byte| TokenId::new(byte as u32)))
1971 {
1972 let mut logits = vec![0.0; FragmentedUtf8Tokenizer::EOS as usize + 1];
1973 processor.mask_logits(&mut logits, &generated).unwrap();
1974 assert!(
1975 logits[token.get() as usize].is_finite(),
1976 "token {} rejected after {:?}",
1977 token.get(),
1978 generated
1979 );
1980 generated.push(token);
1981 }
1982 assert!(processor.is_accepting(&generated).unwrap());
1983 }
1984
1985 #[test]
1986 fn reasoning_delimiter_defers_then_activates_the_grammar() {
1987 let processor = factory()
1988 .create_processor(
1989 &ResponseFormat::JsonObject,
1990 &StructuredOutputStart::AfterDelimiter("</think>".to_string()),
1991 TEST_MAX_OUTPUT_TOKENS,
1992 &HashSet::new(),
1993 &[],
1994 )
1995 .unwrap()
1996 .unwrap();
1997 let mut generated = Vec::new();
1998 let controls = HashSet::from([b'<' as u32, b'>' as u32]);
1999 let mut waiting_logits = vec![0.0; EOS as usize + 1];
2000 let waiting = processor
2001 .mask_logits_with_terminals(&mut waiting_logits, &generated, &HashSet::new(), &controls)
2002 .unwrap();
2003 assert_eq!(waiting.phase, StructuredOutputPhase::WaitingForDelimiter);
2004 assert!(!waiting.accepting);
2005 assert_eq!(waiting.grammar_start_token_index, None);
2006 assert_eq!(waiting.required_delimiter_token_id, Some(b'<' as u32));
2007 assert!(waiting_logits[b'<' as usize].is_finite());
2008 assert!(!waiting_logits[b'>' as usize].is_finite());
2009
2010 let delimiter_prefix = "</think"
2011 .bytes()
2012 .map(|byte| TokenId::new(byte as u32))
2013 .collect::<Vec<_>>();
2014 let mut partial_logits = vec![0.0; EOS as usize + 1];
2015 let partial = processor
2016 .mask_logits_with_terminals(
2017 &mut partial_logits,
2018 &delimiter_prefix,
2019 &HashSet::new(),
2020 &controls,
2021 )
2022 .unwrap();
2023 assert_eq!(partial.phase, StructuredOutputPhase::WaitingForDelimiter);
2024 assert_eq!(partial.grammar_start_token_index, None);
2025 assert_eq!(partial.required_delimiter_token_id, Some(b'>' as u32));
2026 assert!(!partial_logits[b'<' as usize].is_finite());
2027 assert!(partial_logits[b'>' as usize].is_finite());
2028 let partial_progress = processor
2029 .progress_with_terminals(&delimiter_prefix, &HashSet::new())
2030 .unwrap();
2031 assert_eq!(partial_progress.delimiter_token_count, Some(8));
2032 assert_eq!(partial_progress.delimiter_prefix_token_count, 7);
2033
2034 processor.reset().unwrap();
2035
2036 assert_and_append(&processor, &mut generated, "reasoning [is free]</think>");
2037 let mut logits = vec![0.0; EOS as usize + 1];
2038 let active = processor
2039 .mask_logits_with_terminals(&mut logits, &generated, &HashSet::new(), &HashSet::new())
2040 .unwrap();
2041 assert_eq!(active.phase, StructuredOutputPhase::EnforcingGrammar);
2042 assert_eq!(active.grammar_start_token_index, Some(27));
2043 assert!(logits[b'{' as usize].is_finite());
2044 assert!(!logits[b'[' as usize].is_finite());
2045 assert!(!logits[EOS as usize].is_finite());
2046
2047 assert_and_append(&processor, &mut generated, r#"{"ok":true}"#);
2048 assert!(processor.is_accepting(&generated).unwrap());
2049 let progress = processor
2050 .progress_with_terminals(&generated, &HashSet::new())
2051 .unwrap();
2052 assert_eq!(progress.phase, StructuredOutputPhase::EnforcingGrammar);
2053 assert!(progress.accepting);
2054 assert_eq!(progress.generated_token_count, generated.len());
2055 assert!(!progress.boundary_forced);
2056 assert_eq!(progress.reasoning_token_count, Some(19));
2057 }
2058
2059 #[test]
2060 fn reasoning_budget_forces_exact_delimiter_and_preserves_structured_reserve() {
2061 let processor = factory()
2062 .create_processor(
2063 &ResponseFormat::JsonObject,
2064 &StructuredOutputStart::AfterDelimiter("</think>".to_string()),
2065 48,
2066 &HashSet::new(),
2067 &[],
2068 )
2069 .unwrap()
2070 .unwrap();
2071 let mut generated = "reason!!"
2072 .bytes()
2073 .map(|byte| TokenId::new(byte as u32))
2074 .collect::<Vec<_>>();
2075
2076 for expected in "</think>".bytes() {
2077 let mut logits = vec![f32::NEG_INFINITY; EOS as usize + 1];
2078 let outcome = processor
2079 .mask_logits_with_terminals(
2080 &mut logits,
2081 &generated,
2082 &HashSet::new(),
2083 &HashSet::new(),
2084 )
2085 .unwrap();
2086 assert_eq!(outcome.phase, StructuredOutputPhase::ForcingDelimiter);
2087 assert_eq!(outcome.grammar_start_token_index, None);
2088 assert_eq!(outcome.required_delimiter_token_id, Some(expected as u32));
2089 assert_eq!(
2090 logits
2091 .iter()
2092 .enumerate()
2093 .filter(|(_, logit)| logit.is_finite())
2094 .map(|(token, _)| token)
2095 .collect::<Vec<_>>(),
2096 vec![expected as usize]
2097 );
2098 generated.push(TokenId::new(expected as u32));
2099 }
2100
2101 let mut grammar_logits = vec![0.0; EOS as usize + 1];
2102 let outcome = processor
2103 .mask_logits_with_terminals(
2104 &mut grammar_logits,
2105 &generated,
2106 &HashSet::new(),
2107 &HashSet::new(),
2108 )
2109 .unwrap();
2110 assert_eq!(outcome.phase, StructuredOutputPhase::EnforcingGrammar);
2111 assert!(grammar_logits[b'{' as usize].is_finite());
2112 assert!(!grammar_logits[b'[' as usize].is_finite());
2113
2114 let progress = processor
2115 .progress_with_terminals(&generated, &HashSet::new())
2116 .unwrap();
2117 assert_eq!(progress.reasoning_token_count, Some(8));
2118 assert!(progress.boundary_forced);
2119 assert_eq!(
2120 progress.budget,
2121 Some(StructuredOutputBudgetPlan {
2122 total_output_tokens: 48,
2123 reasoning_token_limit: 8,
2124 boundary_token_count: 8,
2125 structured_reserve_tokens: 32,
2126 })
2127 );
2128
2129 processor.reset().unwrap();
2130 let reset_progress = processor
2131 .progress_with_terminals(&[], &HashSet::new())
2132 .unwrap();
2133 assert_eq!(
2134 reset_progress.phase,
2135 StructuredOutputPhase::WaitingForDelimiter
2136 );
2137 assert!(!reset_progress.boundary_forced);
2138 assert_eq!(reset_progress.reasoning_token_count, Some(0));
2139 }
2140
2141 #[test]
2142 fn reasoning_budget_reserves_half_of_a_normal_completion_for_structure() {
2143 assert_eq!(
2144 StructuredOutputBudgetPlan::automatic(1024, 1).unwrap(),
2145 StructuredOutputBudgetPlan {
2146 total_output_tokens: 1024,
2147 reasoning_token_limit: 511,
2148 boundary_token_count: 1,
2149 structured_reserve_tokens: 512,
2150 }
2151 );
2152 }
2153
2154 #[test]
2155 fn reasoning_delimiter_requires_room_beyond_the_boundary() {
2156 let error = factory()
2157 .create_processor(
2158 &ResponseFormat::JsonObject,
2159 &StructuredOutputStart::AfterDelimiter("</think>".to_string()),
2160 8,
2161 &HashSet::new(),
2162 &[],
2163 )
2164 .unwrap_err();
2165 assert!(error
2166 .to_string()
2167 .contains("max_tokens greater than its 8-token delimiter"));
2168 }
2169
2170 #[test]
2171 fn forcing_accounts_for_an_existing_delimiter_prefix() {
2172 let processor = factory()
2173 .create_processor(
2174 &ResponseFormat::JsonObject,
2175 &StructuredOutputStart::AfterDelimiter("</think>".to_string()),
2176 48,
2177 &HashSet::new(),
2178 &[],
2179 )
2180 .unwrap()
2181 .unwrap();
2182 let generated = "reason</"
2183 .bytes()
2184 .map(|byte| TokenId::new(byte as u32))
2185 .collect::<Vec<_>>();
2186 let mut logits = vec![0.0; EOS as usize + 1];
2187
2188 let outcome = processor
2189 .mask_logits_with_terminals(&mut logits, &generated, &HashSet::new(), &HashSet::new())
2190 .unwrap();
2191 assert_eq!(outcome.phase, StructuredOutputPhase::ForcingDelimiter);
2192 assert_eq!(outcome.required_delimiter_token_id, Some(b't' as u32));
2193 let progress = processor
2194 .progress_with_terminals(&generated, &HashSet::new())
2195 .unwrap();
2196 assert_eq!(progress.delimiter_prefix_token_count, 2);
2197 assert_eq!(progress.reasoning_token_count, Some(6));
2198 }
2199
2200 #[test]
2201 fn delimiter_rejects_any_conflicting_stop_condition_up_front() {
2202 let token_error = factory()
2203 .create_processor(
2204 &ResponseFormat::JsonObject,
2205 &StructuredOutputStart::AfterDelimiter("</think>".to_string()),
2206 TEST_MAX_OUTPUT_TOKENS,
2207 &HashSet::from([b'/' as u32]),
2208 &[],
2209 )
2210 .unwrap_err();
2211 assert!(token_error
2212 .to_string()
2213 .contains("conflicts with a stop token"));
2214
2215 let text_error = factory()
2216 .create_processor(
2217 &ResponseFormat::JsonObject,
2218 &StructuredOutputStart::AfterDelimiter("</think>".to_string()),
2219 TEST_MAX_OUTPUT_TOKENS,
2220 &HashSet::new(),
2221 &["think".to_string()],
2222 )
2223 .unwrap_err();
2224 assert!(text_error
2225 .to_string()
2226 .contains("conflicts with stop sequence"));
2227 }
2228}