Skip to main content

ferrum_sampler/
structured_output.rs

1//! Tokenizer-aware hard constraints for structured output.
2//!
3//! A factory owns the tokenizer trie and grammar compiler and is shared by an
4//! engine. Each request gets an independent matcher with no shared mutable
5//! grammar state.
6
7use 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;
29// Structured output is the requested product result; hidden reasoning may use
30// at most the other half of a normal-sized completion budget.
31const 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/// Immutable per-request output budget used when a structured grammar starts
37/// after a reasoning delimiter.
38#[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
90/// Shared, immutable tokenizer and grammar compilation state.
91pub 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    /// Build the tokenizer trie once for this engine.
112    pub fn new(tokenizer: Arc<dyn Tokenizer + Send + Sync>) -> Result<Self> {
113        Self::new_with_model_vocab_size(tokenizer, None)
114    }
115
116    /// Build against the executor's logits width when it is larger than the
117    /// tokenizer base vocabulary (for example added EOS/control tokens).
118    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                    // Keep vocabulary holes out of the trie. The explicit
161                    // eligibility mask below is still required because
162                    // llguidance's wildcard slice represents its root as an
163                    // all-token bitset, including IDs with no trie node.
164                    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    /// Compile one request's grammar while reusing the tokenizer trie.
220    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    /// Compile one request's grammar while reusing the tokenizer trie.
256    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                // Empty reasoning is legal in either mode, so a stop can also
356                // cross directly from the opening header into the closing tag.
357                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    // x-guidance is an llguidance compiler extension, not a JSON Schema
466    // constraint. Keep compiler policy owned by Ferrum so a request cannot
467    // re-enable an unbounded whitespace loop or change JSON separators.
468    JsonCompileOptions {
469        whitespace_flexible: false,
470        ..JsonCompileOptions::default()
471    }
472    .apply_to(&mut schema);
473    Ok(schema)
474}
475
476/// Per-request structured-output parser state.
477pub 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/// Typed phase returned after applying a structured-output constraint.
487#[derive(Debug, Clone, Copy, PartialEq, Eq)]
488pub enum StructuredOutputPhase {
489    WaitingForDelimiter,
490    ForcingDelimiter,
491    EnforcingGrammar,
492}
493
494/// Privacy-safe lexical class for the tail of an incomplete grammar.
495///
496/// This deliberately exposes neither decoded text nor a token history. It is
497/// used only by terminal diagnostics to distinguish an unbounded whitespace,
498/// number, string, or control-token run from a parser activation failure.
499#[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/// Allocation-free hot-path result of one structured-output mask operation.
513#[derive(Debug, Clone, Copy, PartialEq, Eq)]
514pub struct StructuredOutputMaskOutcome {
515    pub phase: StructuredOutputPhase,
516    pub accepting: bool,
517    pub liveness_intervention: bool,
518    /// Generated-token index at which the visible grammar-owned output starts.
519    /// `None` means the processor is still in the hidden pre-grammar domain.
520    /// The execution engine uses this boundary to scope request-local sampling
521    /// history without teaching the grammar about penalty policy.
522    pub grammar_start_token_index: Option<usize>,
523    /// Exact delimiter token authorized for the next sampling step while the
524    /// processor is waiting to activate. The engine uses this typed grant to
525    /// avoid rejecting an intentionally hidden special token during output
526    /// quality filtering.
527    pub required_delimiter_token_id: Option<u32>,
528    /// A native protocol terminal allowed by the complete grammar at this
529    /// step, even though consuming it is necessary for root acceptance.
530    pub grammar_owned_terminal_token_id: Option<u32>,
531}
532
533/// Terminal/debug snapshot that distinguishes activation failures from an
534/// incomplete grammar without retaining generated text.
535#[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    /// Consume newly generated tokens and hard-mask every illegal next token.
596    /// Waiting-for-reasoning mode leaves normal logits untouched until the
597    /// typed delimiter has been observed.
598    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    /// Apply the grammar mask while allowing engine-resolved stop tokens once
604    /// the grammar accepts. Some model templates use an end-of-turn token that
605    /// is not the tokenizer's primary EOS, so it cannot be represented by the
606    /// grammar parser's EOS set alone.
607    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    /// True only when reasoning has closed and the grammar accepts the full
783    /// generated structured value.
784    pub fn is_accepting(&self, generated: &[TokenId]) -> Result<bool> {
785        self.is_accepting_inner(generated, None)
786    }
787
788    /// Completion check that treats an engine-resolved terminal sampled after
789    /// grammar acceptance as framing rather than part of the JSON value.
790    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    /// Inspect the typed activation/grammar state after consuming `generated`.
809    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            // llguidance rollback treats model EOS tokens as zero bytes.
939            // A composed native grammar consumes their actual marker bytes,
940            // so reset from the pristine grammar instead of rolling them back.
941            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    /// Native framing can end inside an ordinary merged token instead of a
964    /// model EOS ID. Only a complete native root authorizes this path.
965    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    /// Byte end of a complete composed wire result, excluding an external
987    /// model EOS. Text-stop handling can compare its cut position with this
988    /// boundary without parsing tool markup as a final JSON value.
989    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    /// Classify a complete non-native result with final-schema precedence.
1014    /// The returned payload preserves bytes from tokens spanning the
1015    /// reasoning/result boundary. Native protocols keep their own parser.
1016    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}