diff --git a/crates/spark-server/src/main_modules/serve_phases/tokenizer_runtime.rs b/crates/spark-server/src/main_modules/serve_phases/tokenizer_runtime.rs index 826144d70..ca44f2651 100644 --- a/crates/spark-server/src/main_modules/serve_phases/tokenizer_runtime.rs +++ b/crates/spark-server/src/main_modules/serve_phases/tokenizer_runtime.rs @@ -287,6 +287,18 @@ pub(crate) fn resolve_tokenizer_runtime( if let Some(tid) = tool_call_end_token { tracing::info!("Tool call end token: {} ({})", tid, tc_end_str); } + // qwen3_coder XML parameter markers for the envelope-stuck guard, from + // THIS tokenizer (they were hard-coded to Qwen3.6 ids; Nemotron's differ). + let tool_markers = + crate::scheduler::tool_markers::ToolMarkers::from_tokenizer(tokenizer.inner()); + match &tool_markers { + Some(m) => tracing::info!("Tool parameter marker classes: {}", m.summary()), + None => tracing::warn!( + "Tool parameter markers unresolved for this tokenizer: parameter-value \ + tracking and the envelope-stuck cap are off (--tool-max-tokens still applies)" + ), + } + crate::scheduler::tool_markers::set(tool_markers); let grammar_engine = { let stop_ids: Vec = eos_tokens.iter().map(|&id| id as i32).collect(); diff --git a/crates/spark-server/src/scheduler/emit_step.rs b/crates/spark-server/src/scheduler/emit_step.rs index ec30b6b76..2a429fcd6 100644 --- a/crates/spark-server/src/scheduler/emit_step.rs +++ b/crates/spark-server/src/scheduler/emit_step.rs @@ -625,8 +625,10 @@ pub enum StartPrefillResult { // - `a.param_body_chars_emitted` ++ per non-close body token // - `a.finished` forced when stuck >MAX_TOOL_BODY_TOKENS // -// Token IDs are Qwen3.6 byte-level BPE (verified via /tokenize 2026-05-25): -// 27 = `<`, 28 = `=`, 29 = `>`, 510 = `` (`>`, +// `>\n`), ends-with `…` /// that is NOT a parameter-value body). Catches a model that opens `` @@ -669,6 +671,13 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { if !a.inside_tool_body { return; } + // Without the served tokenizer's markers the value body cannot be told + // from the envelope, and counting file content as envelope kills large + // writes at MAX_TOOL_BODY_TOKENS. Leave both the tracking and the cap off; + // `--tool-max-tokens` still bounds a runaway envelope. + let Some(m) = super::tool_markers::get() else { + return; + }; // CC6 (2026-06-07): count ONLY envelope tokens — tokens of a // `` VALUE (the file content of a `write` call) are exempt, // so a legitimately large single-file Write streams without tripping the @@ -690,12 +699,6 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { a.finished = true; } - const TOK_LT: u32 = 27; - const TOK_PARAMETER: u32 = 15704; - const TOK_EQ: u32 = 28; - const TOK_GT: u32 = 29; - const TOK_LT_SLASH: u32 = 510; - if a.inside_parameter_body { // P0-1 (2026-07-09): PROVISIONAL close detection. The old code // exited the body on ANY ` { - if tok == TOK_LT_SLASH { + if m.is_lt_slash(tok) { a.param_close_pending = 1; } else { // Any non-close body token advances the counter. The @@ -724,7 +727,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { } } 1 => { - if tok == TOK_PARAMETER { + if m.is_parameter(tok) { a.param_close_pending = 2; } else { // ``), not a close. @@ -734,7 +737,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { } _ => { a.param_close_pending = 0; - if tok == TOK_GT { + if m.is_gt(tok) { // Confirmed `` — exit body. Also reset the // envelope streak: a confirmed close IS forward progress // (the doc comment above always claimed this reset; the @@ -756,7 +759,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { // Not yet inside_parameter_body: scan for `` opener // ending at this `>` (29). Lookback 8 tokens for `[27, 15704, 28]` // signature without an intervening close. - if tok != TOK_GT { + if !m.is_gt(tok) { return; } // Auto-detect whether `tok` is already in output_tokens (caller @@ -776,7 +779,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { let window = &a.output_tokens[start..n_for_lookback]; let mut sig_idx: Option = None; for i in 0..window.len().saturating_sub(2) { - if window[i] == TOK_LT && window[i + 1] == TOK_PARAMETER && window[i + 2] == TOK_EQ { + if m.is_lt(window[i]) && m.is_parameter(window[i + 1]) && m.is_eq_start(window[i + 2]) { sig_idx = Some(i + 3); } } @@ -784,9 +787,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { // Check no intervening `` in the KEY span between // ``. let body_segment = &window[after_eq..]; - let intervening_close = body_segment - .iter() - .any(|&t| t == TOK_LT_SLASH || t == TOK_GT); + let intervening_close = body_segment.iter().any(|&t| m.is_lt_slash(t) || m.is_gt(t)); if !intervening_close { a.inside_parameter_body = true; a.param_body_chars_emitted = 0; diff --git a/crates/spark-server/src/scheduler/mod.rs b/crates/spark-server/src/scheduler/mod.rs index c2ed8f13f..0aab5220a 100644 --- a/crates/spark-server/src/scheduler/mod.rs +++ b/crates/spark-server/src/scheduler/mod.rs @@ -90,6 +90,7 @@ mod teardown; mod test_support; #[cfg(test)] mod think_skip_tests; +pub(crate) mod tool_markers; mod types; #[cfg(test)] mod verify_bias_tests; diff --git a/crates/spark-server/src/scheduler/tool_markers.rs b/crates/spark-server/src/scheduler/tool_markers.rs new file mode 100644 index 000000000..97a0103bd --- /dev/null +++ b/crates/spark-server/src/scheduler/tool_markers.rs @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Which tokens of the served tokenizer can make up the qwen3_coder XML +//! parameter markers (``, ``). +//! +//! `emit_step::update_tool_param_state` tracks whether generation is inside a +//! `` VALUE, whose tokens are exempt from the 1024-token +//! envelope-stuck cap. It used to compare against five exact Qwen3.6 BPE ids +//! (`<` 27, `=` 28, `>` 29, `` 1062, `` with the +//! newline (`>\n` 1561), and both merge punctuation into the close (`.`, and one +//! ending with `, + parameter: Vec, + eq_start: Vec, + gt_start: Vec, + lt_slash_end: Vec, +} + +fn has(set: &[u32], tok: u32) -> bool { + set.binary_search(&tok).is_ok() +} + +impl ToolMarkers { + pub fn is_lt(&self, tok: u32) -> bool { + has(&self.lt, tok) + } + pub fn is_parameter(&self, tok: u32) -> bool { + has(&self.parameter, tok) + } + /// `=`, or `=` merged with the start of the key (`=new`, `=file`). + pub fn is_eq_start(&self, tok: u32) -> bool { + has(&self.eq_start, tok) + } + /// `>`, or `>` merged with what follows (`>\n`, `>\n\n`). + pub fn is_gt(&self, tok: u32) -> bool { + has(&self.gt_start, tok) + } + /// ` bool { + has(&self.lt_slash_end, tok) + } + + /// Build the classes from `(id, decoded text)` for every vocabulary entry. + /// `None` when any class is empty. + pub fn from_vocab<'a>(vocab: impl IntoIterator) -> Option { + let mut m = Self { + lt: Vec::new(), + parameter: Vec::new(), + eq_start: Vec::new(), + gt_start: Vec::new(), + lt_slash_end: Vec::new(), + }; + for (id, text) in vocab { + if text.trim_start() == "<" { + m.lt.push(id); + } + if text == "parameter" { + m.parameter.push(id); + } + if text.starts_with('=') { + m.eq_start.push(id); + } + if text.starts_with('>') { + m.gt_start.push(id); + } + if text.ends_with(" Option { + let n = tok.get_vocab_size(true) as u32; + let texts: Vec<(u32, String)> = (0..n) + .filter_map(|id| tok.decode(&[id], false).ok().map(|t| (id, t))) + .collect(); + Self::from_vocab(texts.iter().map(|(id, t)| (*id, t.as_str()))) + } + + /// Class sizes, for the load log. + pub fn summary(&self) -> String { + format!( + "< x{} parameter x{} =* x{} >* x{} * Self { + Self { + lt: vec![27], + parameter: vec![15704], + eq_start: vec![28], + gt_start: vec![29], + lt_slash_end: vec![510], + } + } +} + +fn legacy() -> Arc { + static LEGACY: std::sync::OnceLock> = std::sync::OnceLock::new(); + LEGACY + .get_or_init(|| Arc::new(ToolMarkers::legacy_qwen())) + .clone() +} + +/// `None` = never installed (the legacy default applies); +/// `Some(None)` = installed, unresolved (tracking off). +type Slot = Option>>; + +#[cfg(not(test))] +static MARKERS: std::sync::RwLock = std::sync::RwLock::new(None); + +#[cfg(test)] +thread_local! { + // Per test thread, so parallel tests installing different markers do not + // observe each other's. + static MARKERS: std::cell::RefCell = const { std::cell::RefCell::new(None) }; +} + +/// Install the classes of the tokenizer just loaded (`None` turns the +/// parameter-body tracking off). +pub fn set(markers: Option) { + let slot: Slot = Some(markers.map(Arc::new)); + #[cfg(not(test))] + { + *MARKERS.write().unwrap_or_else(|e| e.into_inner()) = slot; + } + #[cfg(test)] + MARKERS.with(|m| *m.borrow_mut() = slot); +} + +/// The classes of the tokenizer being served (the legacy Qwen ids until one is +/// installed). One uncontended read lock and an `Arc` clone per call. +pub fn get() -> Option> { + #[cfg(not(test))] + let slot = MARKERS.read().unwrap_or_else(|e| e.into_inner()).clone(); + #[cfg(test)] + let slot = MARKERS.with(|m| m.borrow().clone()); + slot.unwrap_or_else(|| Some(legacy())) +} + +#[cfg(test)] +pub(crate) mod fixtures { + use super::ToolMarkers; + + /// The parts of the Qwen3.6/3.8 vocabulary the markers depend on, with the + /// merges seen in real tool calls (tokenizer.json, 2026-10-18): `=file`, + /// `=re`, `=new`, punctuation-merged closes; `>` and `\n` separate. + /// Merged-form ids >= 900_000 are stand-ins (only their text matters). + pub fn qwen() -> ToolMarkers { + ToolMarkers::from_vocab([ + (27, "<"), + (366, " <"), + (28, "="), + (900_000, "=file"), + (900_001, "=re"), + (900_002, "=new"), + (29, ">"), + (510, "\n` merged (1561), `=new`, `. ToolMarkers { + ToolMarkers::from_vocab([ + (1060, "<"), + (1061, "="), + (900_010, "=new"), + (1062, ">"), + (1561, ">\n"), + (1885, ""), (4, "` closes inside) must stream to `` with the +//! served tokenizer's marker classes, and is what the five hard-coded Qwen ids +//! used to kill at 1024 tokens — on Nemotron for every key, on Qwen for keys +//! whose `=` merges into the key (`=new` `_str`). + +use super::fixtures::{nemotron, qwen}; +use super::{ToolMarkers, set}; +use crate::scheduler::emit_step::update_tool_param_state; +use crate::scheduler::test_support::test_seq; + +const NEMO_TOOL_CALL: u32 = 14; +const NEMO_TOOL_CALL_END: u32 = 15; +const QWEN_TOOL_CALL: u32 = 151_657; +const QWEN_TOOL_CALL_END: u32 = 151_658; + +/// Nemotron ids (tokenizer.json, 2026-10-18): `` +/// `\n\n/tmp/a.html\n\n` +/// `\n` + body + `\n` ``. +fn nemotron_write(body_tokens: usize) -> Vec { + let mut t = vec![NEMO_TOOL_CALL]; + t.extend([1060, 5165, 1061, 16539, 1561]); // \n + t.extend([1060, 31960, 1061, 7368, 1561]); // \n + t.extend([1047, 20307, 22139, 7120, 1010]); // /tmp/a.html\n + t.extend([1885, 31960, 1561]); // \n + t.extend([1060, 31960, 1061, 5431, 1561]); // \n + body(&mut t, body_tokens, [1885, 1950, 1561]); // \n every 50 + t.extend([1885, 31960, 1561]); // \n + t.extend([1885, 5165, 1062]); // + t.push(NEMO_TOOL_CALL_END); + t +} + +/// Qwen ids: `` tokenizes as `<` `parameter` `=new` `_str` +/// `>` (tokenizer.json, 2026-10-18); `=new` is a stand-in id from the fixture. +fn qwen_new_str(body_tokens: usize) -> Vec { + let mut t = vec![QWEN_TOOL_CALL, 198]; + t.extend([27, 1688, 28, 8_888, 29, 198]); // \n (`=` standalone) + t.extend([27, 15704, 900_002, 4_444, 29, 198]); // \n + body(&mut t, body_tokens, [510, 611, 29]); // every 50 + t.extend([198, 510, 15704, 29, 198]); // \n\n + t.extend([510, 1688, 29, 198]); // \n + t.push(QWEN_TOOL_CALL_END); + t +} + +fn body(t: &mut Vec, n: usize, close_tag: [u32; 3]) { + for i in 0..n { + if i % 50 == 49 { + t.extend(close_tag); + } else { + t.push(41_199 + (i as u32 % 7)); + } + } +} + +/// Drive the guard the way `emit_token` does (state update, then push). +fn run( + markers: Option, + open: u32, + close: u32, + tokens: &[u32], +) -> (bool, Option<&'static str>, bool) { + set(markers); + let (mut a, _rx) = test_seq(Vec::new(), 100_000, None, 0); + a.tool_call_start_token = Some(open); + a.tool_call_end_token = Some(close); + // The fixture builds a finished sequence; this one is mid-generation. + a.finished = false; + a.guard_stop = None; + for &tok in tokens { + update_tool_param_state(&mut a, tok, true); + a.output_tokens.push(tok); + if a.finished { + break; + } + } + (a.finished, a.guard_stop, a.inside_tool_body) +} + +fn nemo(markers: Option, tokens: &[u32]) -> (bool, Option<&'static str>, bool) { + run(markers, NEMO_TOOL_CALL, NEMO_TOOL_CALL_END, tokens) +} + +#[test] +fn nemotron_classes_let_a_large_write_finish() { + let (finished, guard, inside) = nemo(Some(nemotron()), &nemotron_write(1500)); + assert!(!finished, "guard fired: {guard:?}"); + assert!(!inside, " should close the body"); +} + +#[test] +fn legacy_qwen_ids_on_a_nemotron_stream_reproduce_the_kill() { + // The old behaviour: Qwen ids never match a Nemotron `` then 1100 envelope tokens + // that never open a parameter value. + let mut t = vec![NEMO_TOOL_CALL]; + t.extend(std::iter::repeat_n(41_199u32, 1100)); + let (finished, guard, _) = nemo(Some(nemotron()), &t); + assert!(finished); + assert_eq!(guard, Some("tool_envelope_stuck")); +}