[Spec] Unify speculative grammar token-accept path in decode processing (#28682)

python/sglang/test/kits/spec_server_kits.py

@@ -620,3 +620,74 @@ class SpecHiddenStatesKit:
             for row in decode_rows:
                 self.assertIsInstance(row, list)
                 self.assertEqual(len(row), hidden_dim)
+
+
+class SpecGrammarKit:
+    """Grammar-constrained structured output under spec decoding.
+
+    Regression for spec verify accepting tokens past grammar termination: the
+    output must be valid JSON with nothing emitted after completion, and the
+    logprob count must match the (truncated) completion-token count.
+    """
+
+    # Override per config if a different schema is desired.
+    grammar_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "name": {"type": "string", "pattern": "^[\\w]+$"},
+                "population": {"type": "integer"},
+                "country": {"type": "string", "pattern": "^[\\w ]+$"},
+                "capital": {"type": "string", "pattern": "^[\\w ]+$"},
+            },
+            "required": ["name", "population", "country", "capital"],
+        }
+    )
+
+    def _generate_grammar(self, return_logprob: bool):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "Here is the information of the capital of France in the JSON format.\n",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 256,
+                    "json_schema": self.grammar_json_schema,
+                },
+                "return_logprob": return_logprob,
+                "logprob_start_len": 0,
+            },
+        )
+        self.assertEqual(response.status_code, 200, response.text)
+        out = response.json()
+        self.assertGreater(
+            out["meta_info"]["spec_verify_ct"],
+            0,
+            "expected spec decoding to run (spec_verify_ct > 0)",
+        )
+        return out
+
+    def test_grammar_structured_output_no_trailing_tokens(self):
+        """Output is valid JSON with nothing emitted past grammar completion."""
+        out = self._generate_grammar(return_logprob=False)
+        text = out["text"]
+        parsed = json.loads(text)
+        for key in ("name", "population", "country", "capital"):
+            self.assertIn(key, parsed)
+        self.assertTrue(
+            text.strip().endswith("}"), f"unexpected trailing tokens: {text!r}"
+        )
+
+    def test_grammar_logprob_count_matches_completion_tokens(self):
+        """Trimmed spec tokens keep logprob count == completion token count."""
+        out = self._generate_grammar(return_logprob=True)
+        meta = out["meta_info"]
+        completion_tokens = meta["completion_tokens"]
+        output_logprobs = meta["output_token_logprobs"]
+        self.assertEqual(
+            len(output_logprobs),
+            completion_tokens,
+            "output logprobs must align with retained (trimmed) tokens: "
+            f"got {len(output_logprobs)} logprobs vs {completion_tokens} completion tokens",
+        )
+        json.loads(out["text"])