fix(skippy): isolate chat grammar during speculative verification (#1172)

* Isolate grammar state during speculative verification

* Align serial grammar verification test

* fix(skippy): trim sampled verification test window

* test(skippy): cover long-context tool verification

crates/skippy-runtime/src/tests.rs

@@ -8,7 +8,13 @@ mod tests {
         GGML_TYPE_F16, ModelInfo, NativeMtpDraft, RuntimeConfig, RuntimeLoadMode, SamplingConfig,
         StageModel, StageSession, Status, TensorRole, format_skippy_error,
     };
-    use std::{env, path::PathBuf};
+    use std::{
+        env,
+        path::PathBuf,
+        time::{Duration, Instant},
+    };
+
+    const TOOL_CALLS_JSON: &str = r#"[{"type":"function","function":{"name":"execute_bash","description":"Run a command.","parameters":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}}]"#;
 
     fn correctness_model() -> Option<PathBuf> {
         env::var_os("SKIPPY_CORRECTNESS_MODEL").map(PathBuf::from)
@@ -47,12 +53,19 @@ mod tests {
     }
 
     fn open_correctness_model(model_path: &PathBuf) -> anyhow::Result<StageModel> {
+        open_correctness_model_with_context(model_path, 256)
+    }
+
+    fn open_correctness_model_with_context(
+        model_path: &PathBuf,
+        ctx_size: u32,
+    ) -> anyhow::Result<StageModel> {
         let layer_end = infer_layer_end(model_path)?;
         let config = RuntimeConfig {
             stage_index: 0,
             layer_start: 0,
             layer_end,
-            ctx_size: 256,
+            ctx_size,
             lane_count: 1,
             n_batch: None,
             n_ubatch: None,
@@ -74,6 +87,13 @@ mod tests {
         StageModel::open(model_path, &config)
     }
 
+    fn tool_call_template_options() -> ChatTemplateJsonOptions {
+        ChatTemplateJsonOptions {
+            tools_json: Some(TOOL_CALLS_JSON.to_string()),
+            ..ChatTemplateJsonOptions::default()
+        }
+    }
+
     #[test]
     fn chat_template_applies_when_model_is_configured() -> anyhow::Result<()> {
         let Some(model_path) = correctness_model() else {
@@ -226,13 +246,7 @@ mod tests {
         let model = open_correctness_model(&model_path)?;
         let rendered = model.apply_chat_template_json(
             r#"[{"role":"user","content":"Call execute_bash."}]"#,
-            ChatTemplateJsonOptions {
-                tools_json: Some(
-                    r#"[{"type":"function","function":{"name":"execute_bash","description":"Run a command.","parameters":{"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}}}]"#
-                        .to_string(),
-                ),
-                ..ChatTemplateJsonOptions::default()
-            },
+            tool_call_template_options(),
         )?;
         let metadata: Value = serde_json::from_str(&rendered.metadata_json)?;
         assert!(
@@ -275,11 +289,14 @@ mod tests {
             prompt_token_count,
             Some(&sampling),
         )?;
-        let serial_predictions = verify_inputs
-            .iter()
-            .copied()
-            .map(|token| serial.decode_step_sampled(token, Some(&sampling)))
-            .collect::<Result<Vec<_>>>()?;
+        let mut serial_predictions = Vec::with_capacity(verify_inputs.len());
+        for (index, token) in verify_inputs.iter().copied().enumerate() {
+            let predicted = serial.decode_step_sampled(token, Some(&sampling))?;
+            serial_predictions.push(predicted);
+            if index + 1 < verify_inputs.len() && predicted != verify_inputs[index + 1] {
+                break;
+            }
+        }
         let serial_token_count = serial.token_count();
         let serial_native_position = serial.native_position()?;
         drop(serial);
@@ -291,15 +308,182 @@ mod tests {
             prompt_token_count,
             Some(&sampling),
         )?;
-        let batched_predictions =
-            batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?;
-
-        assert_eq!(batched_predictions, serial_predictions);
+        let batched_predictions = batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?;
+        assert_eq!(
+            batched_predictions, serial_predictions,
+            "batched verification must stop at the first target mismatch"
+        );
+        batched.trim_session(serial_token_count)?;
         assert_eq!(batched.token_count(), serial_token_count);
         assert_eq!(batched.native_position()?, serial_native_position);
         Ok(())
     }
 
+    #[test]
+    fn long_resident_tool_context_preserves_grammar_and_native_mtp_acceptance() -> anyhow::Result<()>
+    {
+        const MIN_RESIDENT_TOKENS: usize = 8_192;
+        const CONTEXT_SIZE: u32 = 10_240;
+        // Resident prefixes reserve IDs immediately after the active lane IDs.
+        // A single-lane runtime uses `3`, matching the state-handoff harness.
+        const RESIDENT_PREFIX_ID: i32 = 3;
+
+        let Some(model_path) = correctness_model() else {
+            eprintln!("skipping: SKIPPY_CORRECTNESS_MODEL is not set");
+            return Ok(());
+        };
+        let model = open_correctness_model_with_context(&model_path, CONTEXT_SIZE)?;
+
+        let resident_sentence =
+            "The resident tool context records a completed command result cwd workspace. ";
+        let sentence_tokens = model.tokenize(resident_sentence, false)?;
+        assert!(
+            !sentence_tokens.is_empty(),
+            "resident context sentence must tokenize"
+        );
+        let mut resident_sentence_count = MIN_RESIDENT_TOKENS / sentence_tokens.len() + 1;
+        let (rendered, prompt_tokens) = loop {
+            let resident_context = resident_sentence.repeat(resident_sentence_count);
+            let rendered = model.apply_chat_template_json(
+                &format!(
+                    r#"[{{"role":"user","content":"Call execute_bash after this resident context: {resident_context}"}}]"#
+                ),
+                tool_call_template_options(),
+            )?;
+            let prompt_tokens = model.tokenize(&rendered.prompt, true)?;
+            if prompt_tokens.len() >= MIN_RESIDENT_TOKENS {
+                break (rendered, prompt_tokens);
+            }
+            let shortfall = MIN_RESIDENT_TOKENS - prompt_tokens.len();
+            resident_sentence_count +=
+                shortfall * resident_sentence_count / prompt_tokens.len() + 1;
+        };
+        let metadata: Value = serde_json::from_str(&rendered.metadata_json)?;
+        assert_eq!(
+            metadata.get("grammar_lazy").and_then(Value::as_bool),
+            Some(true),
+            "tool grammar must wait for its trigger"
+        );
+
+        assert!(
+            prompt_tokens.len() >= MIN_RESIDENT_TOKENS,
+            "expected at least {MIN_RESIDENT_TOKENS} resident tokens, got {}",
+            prompt_tokens.len()
+        );
+        assert!(
+            prompt_tokens.len() < CONTEXT_SIZE as usize,
+            "resident prompt must leave room for tool sampling"
+        );
+        let prompt_prefix = &prompt_tokens[..prompt_tokens.len() - 1];
+        let prompt_token_count = u64::try_from(prompt_tokens.len())?;
+        let last_prompt_token = *prompt_tokens.last().expect("checked nonempty prompt");
+        let sampling = SamplingConfig {
+            enabled: true,
+            temperature: 0.0,
+            top_p: 0.95,
+            top_k: 40,
+            min_p: 0.05,
+            ..SamplingConfig::default()
+        };
+
+        let mut prefix_owner = model.create_session()?;
+        prefix_owner.prefill_chunked(prompt_prefix)?;
+        prefix_owner.save_prefix(RESIDENT_PREFIX_ID, prompt_prefix.len() as u64)?;
+        drop(prefix_owner);
+
+        let mut verify_inputs = vec![last_prompt_token];
+        verify_inputs.extend(model.tokenize("<tool_call>", false)?);
+        verify_inputs.extend(model.tokenize(
+            "execute_bash<arg_key>command</arg_key><arg_value>pwd</arg_value></tool_call>",
+            false,
+        )?);
+
+        let mut serial =
+            model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?;
+        serial.configure_chat_sampling(
+            &rendered.metadata_json,
+            prompt_token_count,
+            Some(&sampling),
+        )?;
+        let mut serial_predictions = Vec::with_capacity(verify_inputs.len());
+        for (index, token) in verify_inputs.iter().copied().enumerate() {
+            let predicted = serial.decode_step_sampled(token, Some(&sampling))?;
+            serial_predictions.push(predicted);
+            if index + 1 < verify_inputs.len() && predicted != verify_inputs[index + 1] {
+                break;
+            }
+        }
+        let serial_token_count = serial.token_count();
+        let serial_native_position = serial.native_position()?;
+        drop(serial);
+
+        let mut batched =
+            model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?;
+        batched.configure_chat_sampling(
+            &rendered.metadata_json,
+            prompt_token_count,
+            Some(&sampling),
+        )?;
+        let batched_predictions = batched.verify_tokens_sampled(&verify_inputs, Some(&sampling))?;
+        assert_eq!(
+            batched_predictions, serial_predictions,
+            "resident-KV verification must stop at the first tool-grammar mismatch"
+        );
+        batched.trim_session(serial_token_count)?;
+        assert_eq!(batched.token_count(), serial_token_count);
+        assert_eq!(batched.native_position()?, serial_native_position);
+        drop(batched);
+
+        let mut native_mtp =
+            model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?;
+        native_mtp.configure_chat_sampling(
+            &rendered.metadata_json,
+            prompt_token_count,
+            Some(&sampling),
+        )?;
+        let decode_started = Instant::now();
+        let (predicted, draft) =
+            native_mtp.decode_step_sampled_mtp(last_prompt_token, Some(&sampling), 4)?;
+        let resident_decode_elapsed = decode_started.elapsed();
+        assert!(
+            resident_decode_elapsed < Duration::from_secs(30),
+            "sampling after an 8k resident KV prefix took {resident_decode_elapsed:?}"
+        );
+        drop(native_mtp);
+
+        if let Some(draft) = draft {
+            let mut target =
+                model.create_session_from_resident_prefix(RESIDENT_PREFIX_ID, prompt_prefix)?;
+            target.configure_chat_sampling(
+                &rendered.metadata_json,
+                prompt_token_count,
+                Some(&sampling),
+            )?;
+            let mut target_inputs = vec![last_prompt_token, predicted];
+            target_inputs.extend(&draft.token_ids);
+            let target_predictions =
+                target.verify_tokens_sampled(&target_inputs, Some(&sampling))?;
+            assert_eq!(
+                target_predictions.first(),
+                Some(&predicted),
+                "resident target decode must agree with the native-MTP source token"
+            );
+            let accepted_draft_tokens = target_predictions
+                .iter()
+                .skip(1)
+                .zip(&draft.token_ids)
+                .take_while(|(target, draft)| target == draft)
+                .count();
+            assert!(
+                accepted_draft_tokens > 0,
+                "native MTP must accept a draft token after the resident tool context; draft={:?}, target={target_predictions:?}",
+                draft.token_ids
+            );
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn stage_session_exposes_non_frame_native_mtp_decode_api() {
         type DecodeStepSampledMtp = fn(

crates/skippy-server/src/frontend/linear_proposal.rs

@@ -171,7 +171,8 @@ pub struct LinearProposalReceipt {
     pub accepted_proposal_tokens: usize,
     /// Target tokens committed to the response stream.
     pub committed_tokens: Box<[i32]>,
-    /// Prediction observed for every verification row.
+    /// Authoritative prediction prefix through the full-accept boundary or
+    /// first mismatch. Rejected branch-conditioned suffixes are not sampled.
     pub verification_row_predictions: Box<[i32]>,
     /// Prefix length of row predictions that remained canonical.
     pub canonical_prediction_count: usize,

crates/skippy-server/src/frontend/native_mtp/hybrid.rs

@@ -343,17 +343,14 @@ where
     F: FnMut(i32) -> OpenAiResult<bool>,
 {
     let required_predictions = proposal_tokens.len().saturating_add(1);
-    if predicted_tokens.len() < required_predictions {
-        return Err(OpenAiError::backend(format!(
-            "native MTP verify window returned too few tokens: got {} expected {}",
-            predicted_tokens.len(),
-            required_predictions
-        )));
-    }
-
     let mut accepted_proposal_tokens = 0usize;
     for (index, proposal_token) in proposal_tokens.iter().enumerate() {
-        let predicted = predicted_tokens[index];
+        let Some(&predicted) = predicted_tokens.get(index) else {
+            return Err(OpenAiError::backend(format!(
+                "native MTP verify window ended before a decision: got {} predictions after accepting {accepted_proposal_tokens} proposal tokens",
+                predicted_tokens.len()
+            )));
+        };
         let commit_count = index + 1;
         if predicted != *proposal_token {
             return Ok(NativeMtpVerifyWindowDecision {
@@ -373,6 +370,13 @@ where
         }
     }
 
+    if predicted_tokens.len() < required_predictions {
+        return Err(OpenAiError::backend(format!(
+            "native MTP verify window omitted the boundary token after accepting every proposal token: got {} expected {required_predictions}",
+            predicted_tokens.len()
+        )));
+    }
+
     Ok(NativeMtpVerifyWindowDecision {
         accepted_proposal_tokens,
         commit_count: required_predictions.min(max_new_tokens.saturating_sub(generated_len)),
@@ -633,11 +637,26 @@ mod tests {
     #[test]
     fn verify_window_commits_the_target_correction_after_rejection() {
         let decision =
-            classify_native_mtp_verify_window(&[11, 12], &[11, 42, 99], 0, 8, |_| Ok(false))
-                .unwrap();
+            classify_native_mtp_verify_window(&[11, 12], &[11, 42], 0, 8, |_| Ok(false)).unwrap();
 
         assert_eq!(decision.accepted_proposal_tokens, 1);
         assert_eq!(decision.commit_count, 2);
         assert!(decision.rejected);
     }
+
+    #[test]
+    fn verify_window_rejects_a_prefix_that_ends_before_any_decision() {
+        let error =
+            classify_native_mtp_verify_window(&[11, 12], &[11], 0, 8, |_| Ok(false)).unwrap_err();
+
+        assert!(error.to_string().contains("ended before a decision"));
+    }
+
+    #[test]
+    fn verify_window_requires_a_boundary_after_full_acceptance() {
+        let error = classify_native_mtp_verify_window(&[11, 12], &[11, 12], 0, 8, |_| Ok(false))
+            .unwrap_err();
+
+        assert!(error.to_string().contains("omitted the boundary token"));
+    }
 }

third_party/llama.cpp/patches/0066-Stop-chat-grammar-verification-at-the-first-mismatch.patch

@@ -0,0 +1,37 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: Mesh-LLM CI <ci@mesh-llm.local>
+Date: Wed, 5 Aug 2026 15:33:46 +1000
+Subject: [PATCH] Stop chat grammar verification at the first mismatch
+
+---
+ src/skippy.cpp | 12 +++++++++++-
+ 1 file changed, 11 insertions(+), 1 deletion(-)
+
+diff --git a/src/skippy.cpp b/src/skippy.cpp
+index 6942f08a4..cfc0bac8f 100644
+--- a/src/skippy.cpp
++++ b/src/skippy.cpp
+@@ -6631,12 +6631,22 @@ enum skippy_status skippy_verify_tokens_frame_sampled(
+             return SKIPPY_STATUS_RUNTIME_ERROR;
+         }
+         session->token_history.resize(history_size_before_verification);
++        const bool stop_after_first_mismatch = session->grammar_sampler != nullptr;
++        bool proposal_mismatched = false;
+         for (size_t i = 0; i < token_count; ++i) {
+             skippy_record_tokens(session, &token_ids[i], 1);
+             const int32_t logits_index = static_cast<int32_t>(i);
+             output_tokens[i] = skippy_sample_token_ith(session, sampling, logits_index);
++            if (stop_after_first_mismatch && i + 1 < token_count && output_tokens[i] != token_ids[i + 1]) {
++                // Later rows are conditioned on a rejected token and are not
++                // authoritative predictions. Stop advancing the sampler and
++                // grammar state through that rejected suffix.
++                proposal_mismatched = true;
++                *out_token_count = i + 1;
++                break;
++            }
+         }
+-        if (out_mtp_draft != nullptr && token_count > 0) {
++        if (out_mtp_draft != nullptr && token_count > 0 && !proposal_mismatched) {
+             status = skippy_mtp_propose_next(
+                     session,
+                     output_tokens[token_count - 1],