From 3ed23a1cdf865e643df683eaec8af2e6004e0ec8 Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:05:24 -0400 Subject: [PATCH 1/5] scheduler: resolve the tool-parameter marker ids from the served tokenizer (Nemotron large tool calls were cut at 1024 tokens) update_tool_param_state tracks whether generation is inside a qwen3_coder value, whose tokens are exempt from the 1024-token envelope-stuck cap. It compared against hard-coded Qwen3.6 ids (< 27, = 28, > 29, \n 1561 / 1885 / 31960; the Qwen ids are , , ... there), so a value was never recognised: every file content token counted as envelope, any argument past 1024 tokens tripped tool_envelope_stuck, and the client got a salvaged fragment. Observed in an agent session: large create_file writes landing as the 11 bytes '\n forms, and read by the guard. A tokenizer that does not encode a marker as one token leaves the tracking and the cap off instead of misfiring; --tool-max-tokens still bounds a runaway envelope. Qwen ids resolve to the same values as before, so Qwen behaviour is unchanged. Tests: the resolver on Qwen and Nemotron tables; the real guard on a Nemotron write with a 1500-token content (with closes) finishes with the Nemotron markers, trips with the Qwen ids (the old behaviour), stays off with none, and a never-closed envelope is still cut. Co-Authored-By: Claude Opus 5.5 --- .../serve_phases/tokenizer_runtime.rs | 13 ++ .../spark-server/src/scheduler/emit_step.rs | 31 ++-- crates/spark-server/src/scheduler/mod.rs | 1 + .../src/scheduler/tool_markers.rs | 146 ++++++++++++++++++ .../src/scheduler/tool_markers_emit_tests.rs | 92 +++++++++++ 5 files changed, 269 insertions(+), 14 deletions(-) create mode 100644 crates/spark-server/src/scheduler/tool_markers.rs create mode 100644 crates/spark-server/src/scheduler/tool_markers_emit_tests.rs 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..28cfc705b 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,19 @@ 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::resolve(|s| { + tokenizer.encode(s).ok() + }); + match tool_markers { + Some(m) => tracing::info!("Tool parameter markers: {m:?}"), + 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..f542da30d 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 = ``=29 ``=1062 (and `>\n`=1561) `…` /// 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 tok == m.lt_slash { 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 tok == m.parameter { 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 window[i] == m.lt && window[i + 1] == m.parameter && window[i + 2] == m.eq { sig_idx = Some(i + 3); } } @@ -786,7 +789,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { let body_segment = &window[after_eq..]; let intervening_close = body_segment .iter() - .any(|&t| t == TOK_LT_SLASH || t == TOK_GT); + .any(|&t| t == m.lt_slash || 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..c1a7de47b 100644 --- a/crates/spark-server/src/scheduler/mod.rs +++ b/crates/spark-server/src/scheduler/mod.rs @@ -75,6 +75,7 @@ mod prefill_preempt; mod repetition; mod rollback; mod sample_step; +pub(crate) mod tool_markers; pub mod sched_ctx; mod shared_prefix; #[cfg(test)] 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..dacb79e02 --- /dev/null +++ b/crates/spark-server/src/scheduler/tool_markers.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Token ids of the qwen3_coder XML parameter markers (`<`, `=`, `>`, `` VALUE, whose tokens are exempt from the 1024-token +//! envelope-stuck cap. It used to compare against Qwen3.6 BPE ids +//! (27/28/29/510/15704). Nemotron 3.5 Super VL serves the same XML tool format +//! with its own tokenizer (`<`=1060, `=`=1061, `>`=1062, `\n` merged into one token, 1561), so the value +//! body was never recognised: every file-content token counted as envelope, +//! any argument past 1024 tokens tripped `tool_envelope_stuck`, and the +//! client got a salvaged fragment (`` alone, plus `>\n` / `>\n\n` where the tokenizer merges them into + /// one token (Nemotron does; Qwen does not). + pub gt: [Option; 3], + pub lt_slash: u32, + pub parameter: u32, +} + +impl ToolMarkers { + pub fn is_gt(&self, tok: u32) -> bool { + self.gt.contains(&Some(tok)) + } + + /// Resolve from an encoder (`text -> ids`, no special tokens). `None` when + /// any required marker is not a single token. + pub fn resolve(encode: impl Fn(&str) -> Option>) -> Option { + let one = |s: &str| encode(s).filter(|ids| ids.len() == 1).map(|ids| ids[0]); + Some(Self { + lt: one("<")?, + eq: one("=")?, + gt: [one(">"), one(">\n"), one(">\n\n")], + lt_slash: one("> = RwLock::new(None); + +/// Install the markers of the tokenizer just loaded (`None` turns the +/// parameter-body tracking off). +pub fn set(markers: Option) { + *MARKERS.write().unwrap_or_else(|e| e.into_inner()) = markers; +} + +/// The markers of the tokenizer being served. +pub fn get() -> Option { + *MARKERS.read().unwrap_or_else(|e| e.into_inner()) +} + +/// Serialises the tests (here and in `emit_step`) that install markers into +/// the process-wide slot, so parallel tests do not observe each other's. +#[cfg(test)] +pub(crate) static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + +#[cfg(test)] +pub(crate) mod fixtures { + use super::ToolMarkers; + + /// Qwen3.6 / 3.8 byte-level BPE (verified with tokenizer.json 2026-10-18). + pub const QWEN: ToolMarkers = ToolMarkers { + lt: 27, + eq: 28, + gt: [Some(29), None, None], + lt_slash: 510, + parameter: 15704, + }; + /// Nemotron 3.5 Super VL (tokenizer.json of + /// nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4, 2026-10-18). + pub const NEMOTRON: ToolMarkers = ToolMarkers { + lt: 1060, + eq: 1061, + gt: [Some(1062), Some(1561), None], + lt_slash: 1885, + parameter: 31960, + }; +} + +#[cfg(test)] +mod tests { + use super::fixtures::{NEMOTRON, QWEN}; + use super::*; + + fn enc(table: &'static [(&'static str, &'static [u32])]) -> impl Fn(&str) -> Option> { + move |s| table.iter().find(|(k, _)| *k == s).map(|(_, v)| v.to_vec()) + } + + #[test] + fn resolves_qwen_and_nemotron_markers() { + let qwen = enc(&[ + ("<", &[27]), + ("=", &[28]), + (">", &[29]), + (">\n", &[29, 198]), + (">\n\n", &[29, 271]), + ("", &[1062]), + (">\n", &[1561]), + (">\n\n", &[1062, 1010, 1010]), + ("", &[3]), + ("` closes inside) must +//! stream to `` with the served tokenizer's markers, and is what +//! the Qwen-hardcoded markers used to kill at 1024 tokens. + +use super::fixtures::{NEMOTRON, QWEN}; +use super::{TEST_LOCK, 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; + +/// `` `\n` `\n/tmp/a.html\n\n` +/// `\n` + body + `\n` ``, +/// in Nemotron ids (tokenizer.json, 2026-10-18). +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 + for i in 0..body_tokens { + // Mostly prose ids, with a `\n` (\n) every 50 tokens. + if i % 50 == 49 { + t.extend([1885, 1950, 1561]); + } else { + t.push(41199 + (i as u32 % 7)); + } + } + t.extend([1885, 31960, 1561]); // \n + t.extend([1885, 5165, 1062]); // + t.push(NEMO_TOOL_CALL_END); + t +} + +/// Drive the guard the way `emit_token` does (state update, then push). +fn run(markers: Option, 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(NEMO_TOOL_CALL); + a.tool_call_end_token = Some(NEMO_TOOL_CALL_END); + for &tok in tokens { + update_tool_param_state(&mut a, tok, true); + a.output_tokens.push(tok); + if a.finished { + break; + } + } + set(None); + (a.finished, a.guard_stop, a.inside_tool_body) +} + +#[test] +fn nemotron_markers_let_a_large_write_finish() { + let _g = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let (finished, guard, inside) = run(Some(NEMOTRON), &nemotron_write(1500)); + assert!(!finished, "guard fired: {guard:?}"); + assert_eq!(guard, None); + assert!(!inside, " should close the body"); +} + +#[test] +fn qwen_ids_on_a_nemotron_stream_reproduce_the_kill() { + // The pre-fix behaviour: Qwen ids never match a Nemotron `` then 1100 envelope tokens + // that never open a parameter value. + let _g = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let mut t = vec![NEMO_TOOL_CALL]; + t.extend(std::iter::repeat_n(41199u32, 1100)); + let (finished, guard, _) = run(Some(NEMOTRON), &t); + assert!(finished); + assert_eq!(guard, Some("tool_envelope_stuck")); +} From 217776ccc4b77dee7b2e365d2cdf9d57c31b281a Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:06:48 -0400 Subject: [PATCH 2/5] scheduler: rustfmt the tool-marker wiring Co-Authored-By: Claude Opus 5.5 --- .../src/main_modules/serve_phases/tokenizer_runtime.rs | 5 ++--- crates/spark-server/src/scheduler/emit_step.rs | 4 +--- crates/spark-server/src/scheduler/mod.rs | 2 +- 3 files changed, 4 insertions(+), 7 deletions(-) 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 28cfc705b..b87603fb1 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 @@ -289,9 +289,8 @@ pub(crate) fn resolve_tokenizer_runtime( } // 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::resolve(|s| { - tokenizer.encode(s).ok() - }); + let tool_markers = + crate::scheduler::tool_markers::ToolMarkers::resolve(|s| tokenizer.encode(s).ok()); match tool_markers { Some(m) => tracing::info!("Tool parameter markers: {m:?}"), None => tracing::warn!( diff --git a/crates/spark-server/src/scheduler/emit_step.rs b/crates/spark-server/src/scheduler/emit_step.rs index f542da30d..56b763cce 100644 --- a/crates/spark-server/src/scheduler/emit_step.rs +++ b/crates/spark-server/src/scheduler/emit_step.rs @@ -787,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 == m.lt_slash || m.is_gt(t)); + let intervening_close = body_segment.iter().any(|&t| t == m.lt_slash || 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 c1a7de47b..0aab5220a 100644 --- a/crates/spark-server/src/scheduler/mod.rs +++ b/crates/spark-server/src/scheduler/mod.rs @@ -75,7 +75,6 @@ mod prefill_preempt; mod repetition; mod rollback; mod sample_step; -pub(crate) mod tool_markers; pub mod sched_ctx; mod shared_prefix; #[cfg(test)] @@ -91,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; From b1a9621fa2e974ad13a6c4d53f776a23b879147a Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:08:48 -0400 Subject: [PATCH 3/5] scheduler: tool markers default to the Qwen ids until a tokenizer installs its own; per-thread in tests The first version started the slot empty, which turned the parameter tracking off for any caller that never loads a tokenizer, including the existing think_commit span test that replays a Qwen tool call. The slot now starts at the Qwen ids the guard was hard-coded to (every serve path installs the served tokenizer's markers in tokenizer_runtime). Under cfg(test) the slot is thread-local, so tests installing other markers cannot race the ones that rely on the default. The emit tests also reset the fixture's finished flag; the shared fixture builds a finished sequence. Co-Authored-By: Claude Opus 5.5 --- .../src/scheduler/tool_markers.rs | 50 +++++++++++++------ .../src/scheduler/tool_markers_emit_tests.rs | 11 ++-- 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/crates/spark-server/src/scheduler/tool_markers.rs b/crates/spark-server/src/scheduler/tool_markers.rs index dacb79e02..accf6691e 100644 --- a/crates/spark-server/src/scheduler/tool_markers.rs +++ b/crates/spark-server/src/scheduler/tool_markers.rs @@ -18,6 +18,7 @@ //! token resolves to `None`, and the guard then stays off rather than //! misfiring: `--tool-max-tokens` still bounds a runaway envelope. +#[cfg(not(test))] use std::sync::RwLock; /// The marker ids for one tokenizer. @@ -52,36 +53,55 @@ impl ToolMarkers { } } -static MARKERS: RwLock> = RwLock::new(None); +/// Qwen3.6 / 3.8 byte-level BPE ids: the values the guard was hard-coded to, +/// kept as the default until a tokenizer installs its own (every serve path +/// does, in `tokenizer_runtime`), so nothing changes for a caller that never +/// loads one. +pub const QWEN: ToolMarkers = ToolMarkers { + lt: 27, + eq: 28, + gt: [Some(29), None, None], + lt_slash: 510, + parameter: 15704, +}; + +#[cfg(not(test))] +static MARKERS: RwLock> = RwLock::new(Some(QWEN)); + +#[cfg(test)] +thread_local! { + // Per test thread, so parallel tests installing different markers do not + // observe each other's. + static MARKERS: std::cell::Cell> = const { std::cell::Cell::new(Some(QWEN)) }; +} /// Install the markers of the tokenizer just loaded (`None` turns the /// parameter-body tracking off). pub fn set(markers: Option) { - *MARKERS.write().unwrap_or_else(|e| e.into_inner()) = markers; + #[cfg(not(test))] + { + *MARKERS.write().unwrap_or_else(|e| e.into_inner()) = markers; + } + #[cfg(test)] + MARKERS.with(|m| m.set(markers)); } /// The markers of the tokenizer being served. pub fn get() -> Option { - *MARKERS.read().unwrap_or_else(|e| e.into_inner()) + #[cfg(not(test))] + { + *MARKERS.read().unwrap_or_else(|e| e.into_inner()) + } + #[cfg(test)] + MARKERS.with(|m| m.get()) } -/// Serialises the tests (here and in `emit_step`) that install markers into -/// the process-wide slot, so parallel tests do not observe each other's. -#[cfg(test)] -pub(crate) static TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - #[cfg(test)] pub(crate) mod fixtures { use super::ToolMarkers; /// Qwen3.6 / 3.8 byte-level BPE (verified with tokenizer.json 2026-10-18). - pub const QWEN: ToolMarkers = ToolMarkers { - lt: 27, - eq: 28, - gt: [Some(29), None, None], - lt_slash: 510, - parameter: 15704, - }; + pub use super::QWEN; /// Nemotron 3.5 Super VL (tokenizer.json of /// nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4, 2026-10-18). pub const NEMOTRON: ToolMarkers = ToolMarkers { diff --git a/crates/spark-server/src/scheduler/tool_markers_emit_tests.rs b/crates/spark-server/src/scheduler/tool_markers_emit_tests.rs index 5360377ad..001bec41a 100644 --- a/crates/spark-server/src/scheduler/tool_markers_emit_tests.rs +++ b/crates/spark-server/src/scheduler/tool_markers_emit_tests.rs @@ -5,7 +5,7 @@ //! the Qwen-hardcoded markers used to kill at 1024 tokens. use super::fixtures::{NEMOTRON, QWEN}; -use super::{TEST_LOCK, ToolMarkers, set}; +use super::{ToolMarkers, set}; use crate::scheduler::emit_step::update_tool_param_state; use crate::scheduler::test_support::test_seq; @@ -42,6 +42,9 @@ fn run(markers: Option, tokens: &[u32]) -> (bool, Option<&'static s let (mut a, _rx) = test_seq(Vec::new(), 100_000, None, 0); a.tool_call_start_token = Some(NEMO_TOOL_CALL); a.tool_call_end_token = Some(NEMO_TOOL_CALL_END); + // 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); @@ -49,13 +52,12 @@ fn run(markers: Option, tokens: &[u32]) -> (bool, Option<&'static s break; } } - set(None); + set(Some(super::QWEN)); (a.finished, a.guard_stop, a.inside_tool_body) } #[test] fn nemotron_markers_let_a_large_write_finish() { - let _g = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let (finished, guard, inside) = run(Some(NEMOTRON), &nemotron_write(1500)); assert!(!finished, "guard fired: {guard:?}"); assert_eq!(guard, None); @@ -66,7 +68,6 @@ fn nemotron_markers_let_a_large_write_finish() { fn qwen_ids_on_a_nemotron_stream_reproduce_the_kill() { // The pre-fix behaviour: Qwen ids never match a Nemotron `` then 1100 envelope tokens // that never open a parameter value. - let _g = TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner()); let mut t = vec![NEMO_TOOL_CALL]; t.extend(std::iter::repeat_n(41199u32, 1100)); let (finished, guard, _) = run(Some(NEMOTRON), &t); From d89c8b03d1d2ab8444212497913f97dcc996458f Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:22:26 -0400 Subject: [PATCH 4/5] scheduler: tool-parameter markers are token classes from the vocabulary, not five ids An in-stream tokenization battery (both tokenizers, real tokenizer.json) showed exact-id matching misses openers even on Qwen: BPE merges the '=' into the key's first sub-word for common keys ('=file' '_path', '=re' 'placements', '=new' '_str'), so those values were never recognised and their content counted toward the 1024-token envelope cap. Nemotron merges '>' with the newline ('>\n') and both merge punctuation into the close ('.', ends-with ' off; a never-closed envelope is still cut. Co-Authored-By: Claude Opus 5.5 --- .../serve_phases/tokenizer_runtime.rs | 6 +- .../spark-server/src/scheduler/emit_step.rs | 18 +- .../src/scheduler/tool_markers.rs | 313 ++++++++++++------ .../src/scheduler/tool_markers_emit_tests.rs | 103 ++++-- 4 files changed, 294 insertions(+), 146 deletions(-) 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 b87603fb1..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 @@ -290,9 +290,9 @@ pub(crate) fn resolve_tokenizer_runtime( // 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::resolve(|s| tokenizer.encode(s).ok()); - match tool_markers { - Some(m) => tracing::info!("Tool parameter markers: {m:?}"), + 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)" diff --git a/crates/spark-server/src/scheduler/emit_step.rs b/crates/spark-server/src/scheduler/emit_step.rs index 56b763cce..2bcaa896c 100644 --- a/crates/spark-server/src/scheduler/emit_step.rs +++ b/crates/spark-server/src/scheduler/emit_step.rs @@ -625,10 +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 the SERVED tokenizer's (`tool_markers`), resolved at load: -// Qwen3.6/3.8 `<`=27 `=`=28 `>`=29 ``=1062 (and `>\n`=1561) `` (`>`, +// `>\n`), ends-with `…` /// that is NOT a parameter-value body). Catches a model that opens `` @@ -716,7 +716,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { // next exact close still exits. match a.param_close_pending { 0 => { - if tok == m.lt_slash { + if m.is_lt_slash(tok) { a.param_close_pending = 1; } else { // Any non-close body token advances the counter. The @@ -727,7 +727,7 @@ pub fn update_tool_param_state(a: &mut ActiveSeq, tok: u32, quiet: bool) { } } 1 => { - if tok == m.parameter { + if m.is_parameter(tok) { a.param_close_pending = 2; } else { // ``), not a close. @@ -779,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] == m.lt && window[i + 1] == m.parameter && window[i + 2] == m.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); } } @@ -787,7 +787,9 @@ 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 == m.lt_slash || m.is_gt(t)); + 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/tool_markers.rs b/crates/spark-server/src/scheduler/tool_markers.rs index accf6691e..7528eba52 100644 --- a/crates/spark-server/src/scheduler/tool_markers.rs +++ b/crates/spark-server/src/scheduler/tool_markers.rs @@ -1,163 +1,266 @@ // SPDX-License-Identifier: AGPL-3.0-only -//! Token ids of 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 Qwen3.6 BPE ids -//! (27/28/29/510/15704). Nemotron 3.5 Super VL serves the same XML tool format -//! with its own tokenizer (`<`=1060, `=`=1061, `>`=1062, `\n` merged into one token, 1561), so the value -//! body was never recognised: every file-content token counted as envelope, -//! any argument past 1024 tokens tripped `tool_envelope_stuck`, and the -//! client got a salvaged fragment (`` 29, `` 1062, `` with the +//! newline (`>\n` 1561), and both merge punctuation into the close (`.`, and one +//! ending with `` alone, plus `>\n` / `>\n\n` where the tokenizer merges them into - /// one token (Nemotron does; Qwen does not). - pub gt: [Option; 3], - pub lt_slash: u32, - pub parameter: u32, + lt: Vec, + 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 { - self.gt.contains(&Some(tok)) + 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{} * ids`, no special tokens). `None` when - /// any required marker is not a single token. - pub fn resolve(encode: impl Fn(&str) -> Option>) -> Option { - let one = |s: &str| encode(s).filter(|ids| ids.len() == 1).map(|ids| ids[0]); - Some(Self { - lt: one("<")?, - eq: one("=")?, - gt: [one(">"), one(">\n"), one(">\n\n")], - lt_slash: one(" Self { + Self { + lt: vec![27], + parameter: vec![15704], + eq_start: vec![28], + gt_start: vec![29], + lt_slash_end: vec![510], + } } } -/// Qwen3.6 / 3.8 byte-level BPE ids: the values the guard was hard-coded to, -/// kept as the default until a tokenizer installs its own (every serve path -/// does, in `tokenizer_runtime`), so nothing changes for a caller that never -/// loads one. -pub const QWEN: ToolMarkers = ToolMarkers { - lt: 27, - eq: 28, - gt: [Some(29), None, None], - lt_slash: 510, - parameter: 15704, -}; +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: RwLock> = RwLock::new(Some(QWEN)); +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::Cell> = const { std::cell::Cell::new(Some(QWEN)) }; + static MARKERS: std::cell::RefCell = const { std::cell::RefCell::new(None) }; } -/// Install the markers of the tokenizer just loaded (`None` turns the +/// 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()) = markers; + *MARKERS.write().unwrap_or_else(|e| e.into_inner()) = slot; } #[cfg(test)] - MARKERS.with(|m| m.set(markers)); + MARKERS.with(|m| *m.borrow_mut() = slot); } -/// The markers of the tokenizer being served. -pub fn get() -> Option { +/// 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))] - { - *MARKERS.read().unwrap_or_else(|e| e.into_inner()) - } + let slot = MARKERS.read().unwrap_or_else(|e| e.into_inner()).clone(); #[cfg(test)] - MARKERS.with(|m| m.get()) + let slot = MARKERS.with(|m| m.borrow().clone()); + slot.unwrap_or_else(|| Some(legacy())) } #[cfg(test)] pub(crate) mod fixtures { use super::ToolMarkers; - /// Qwen3.6 / 3.8 byte-level BPE (verified with tokenizer.json 2026-10-18). - pub use super::QWEN; - /// Nemotron 3.5 Super VL (tokenizer.json of - /// nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4, 2026-10-18). - pub const NEMOTRON: ToolMarkers = ToolMarkers { - lt: 1060, - eq: 1061, - gt: [Some(1062), Some(1561), None], - lt_slash: 1885, - parameter: 31960, - }; + /// 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, " impl Fn(&str) -> Option> { - move |s| table.iter().find(|(k, _)| *k == s).map(|(_, v)| v.to_vec()) + #[test] + fn classes_cover_the_merged_forms() { + let q = qwen(); + assert!(q.is_lt(27) && q.is_lt(366) && !q.is_lt(510)); + assert!(q.is_eq_start(28) && q.is_eq_start(900_000) && q.is_eq_start(900_002)); + assert!(q.is_gt(29) && !q.is_gt(198)); + assert!(q.is_lt_slash(510) && q.is_lt_slash(900_003) && q.is_lt_slash(900_004)); + assert!(q.is_parameter(15704)); + let n = nemotron(); + assert!(n.is_gt(1062) && n.is_gt(1561) && !n.is_gt(29)); + assert!(n.is_eq_start(900_010) && n.is_lt_slash(900_011)); } #[test] - fn resolves_qwen_and_nemotron_markers() { - let qwen = enc(&[ - ("<", &[27]), - ("=", &[28]), - (">", &[29]), - (">\n", &[29, 198]), - (">\n\n", &[29, 271]), - ("", &[1062]), - (">\n", &[1561]), - (">\n\n", &[1062, 1010, 1010]), - (""), (4, "", &[3]), - ("` closes inside) must -//! stream to `` with the served tokenizer's markers, and is what -//! the Qwen-hardcoded markers used to kill at 1024 tokens. +//! The envelope-stuck guard on real tool-call token streams: a value of 1500 +//! tokens (with `` 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::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; -/// `` `\n` `\n/tmp/a.html\n\n` -/// `\n` + body + `\n` ``, -/// in Nemotron ids (tokenizer.json, 2026-10-18). +/// 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 @@ -22,26 +25,47 @@ fn nemotron_write(body_tokens: usize) -> Vec { 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 - for i in 0..body_tokens { - // Mostly prose ids, with a `\n` (\n) every 50 tokens. - if i % 50 == 49 { - t.extend([1885, 1950, 1561]); - } else { - t.push(41199 + (i as u32 % 7)); - } - } + 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, tokens: &[u32]) -> (bool, Option<&'static str>, bool) { +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(NEMO_TOOL_CALL); - a.tool_call_end_token = Some(NEMO_TOOL_CALL_END); + 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; @@ -52,30 +76,49 @@ fn run(markers: Option, tokens: &[u32]) -> (bool, Option<&'static s break; } } - set(Some(super::QWEN)); (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_markers_let_a_large_write_finish() { - let (finished, guard, inside) = run(Some(NEMOTRON), &nemotron_write(1500)); +fn nemotron_classes_let_a_large_write_finish() { + let (finished, guard, inside) = nemo(Some(nemotron()), &nemotron_write(1500)); assert!(!finished, "guard fired: {guard:?}"); - assert_eq!(guard, None); assert!(!inside, " should close the body"); } #[test] -fn qwen_ids_on_a_nemotron_stream_reproduce_the_kill() { - // The pre-fix 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(41199u32, 1100)); - let (finished, guard, _) = run(Some(NEMOTRON), &t); + 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")); } From f69e599fa86fd886b71fed90a2acce060cd6f91a Mon Sep 17 00:00:00 2001 From: Azeez Ishaqui <95100110+AzeezIsh@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:24:16 -0400 Subject: [PATCH 5/5] scheduler: rustfmt the marker classes Co-Authored-By: Claude Opus 5.5 --- crates/spark-server/src/scheduler/emit_step.rs | 4 +--- crates/spark-server/src/scheduler/tool_markers.rs | 12 +++++++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/crates/spark-server/src/scheduler/emit_step.rs b/crates/spark-server/src/scheduler/emit_step.rs index 2bcaa896c..2a429fcd6 100644 --- a/crates/spark-server/src/scheduler/emit_step.rs +++ b/crates/spark-server/src/scheduler/emit_step.rs @@ -787,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| m.is_lt_slash(t) || m.is_gt(t)); + 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/tool_markers.rs b/crates/spark-server/src/scheduler/tool_markers.rs index 7528eba52..97a0103bd 100644 --- a/crates/spark-server/src/scheduler/tool_markers.rs +++ b/crates/spark-server/src/scheduler/tool_markers.rs @@ -101,9 +101,15 @@ impl ToolMarkers { v.sort_unstable(); v.dedup(); } - let complete = [&m.lt, &m.parameter, &m.eq_start, &m.gt_start, &m.lt_slash_end] - .iter() - .all(|v| !v.is_empty()); + let complete = [ + &m.lt, + &m.parameter, + &m.eq_start, + &m.gt_start, + &m.lt_slash_end, + ] + .iter() + .all(|v| !v.is_empty()); complete.then_some(m) }