[Rust] Keep sampling and scheduler wire schemas in sync (#37222)

rust/sglang-server/src/message/sampling.rs

@@ -28,6 +28,32 @@ const MAX_STOP_REGEX_LEN: usize = 256;
 /// (`re._MAXCACHE`), so past that every pattern recompiles on every decode step.
 const MAX_STOP_REGEX_COUNT: usize = 32;
 
+/// JSON values accepted by Python's `CustomParamValue`: a scalar, a list of
+/// scalars, or a string-keyed object whose values are scalars.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum CustomParamValue {
+    Null(()),
+    Bool(bool),
+    Signed(i64),
+    Unsigned(u64),
+    Float(f64),
+    String(String),
+    List(Vec<JsonScalar>),
+    Object(BTreeMap<String, JsonScalar>),
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum JsonScalar {
+    Null(()),
+    Bool(bool),
+    Signed(i64),
+    Unsigned(u64),
+    Float(f64),
+    String(String),
+}
+
 /// One module per field default, each exposing the two hooks serde needs under
 /// one name: `default` (key absent) and `deserialize` (key present — including
 /// an explicit `null`, which Python's `__post_init__` maps back to the default:
@@ -180,10 +206,10 @@ pub struct SamplingParams {
     pub logit_bias: Option<BTreeMap<String, f64>>,
     #[serde(default)]
     pub sampling_seed: Option<i64>,
-    /// Opaque JSON object forwarded to a custom logit processor. Python types it
-    /// as `Dict[str, JsonScalar | list | dict]`; it is never inspected here.
+    /// JSON object forwarded to a custom logit processor. Its values match
+    /// Python's `CustomParamValue` exactly.
     #[serde(default)]
-    pub custom_params: Option<serde_json::Value>,
+    pub custom_params: Option<BTreeMap<String, CustomParamValue>>,
 
     // Normalized internal fields.
     //
@@ -329,6 +355,16 @@ impl SamplingParams {
         if self.top_k == -1 {
             self.top_k = TOP_K_ALL; // -1 disables top_k → whole vocabulary
         }
+        for constraint in [
+            &mut self.json_schema,
+            &mut self.regex,
+            &mut self.ebnf,
+            &mut self.structural_tag,
+        ] {
+            if constraint.as_deref() == Some("") {
+                *constraint = None;
+            }
+        }
     }
 
     /// Python `normalize(tokenizer)`: size the stop match windows, reject
@@ -402,6 +438,13 @@ impl SamplingParams {
     /// Python `verify(vocab_size)` — the same ranges, messages and mutual
     /// exclusions, plus the rust-server `n == 1` restriction.
     fn verify(&self, vocab_size: u64) -> Result<(), Error> {
+        if let Some(beam_width) = self.beam_width
+            && beam_width < 1
+        {
+            return Err(bad(format!(
+                "beam_width must be at least 1, got {beam_width}."
+            )));
+        }
         if !self.temperature.is_finite() || self.temperature < 0.0 {
             return Err(bad(format!(
                 "temperature must be a non-negative finite number, got {}",
@@ -477,13 +520,18 @@ impl SamplingParams {
             }
         }
         // Grammars are mutually exclusive.
-        let grammars = [&self.json_schema, &self.regex, &self.ebnf]
-            .iter()
-            .filter(|g| g.is_some())
-            .count();
+        let grammars = [
+            &self.json_schema,
+            &self.regex,
+            &self.ebnf,
+            &self.structural_tag,
+        ]
+        .iter()
+        .filter(|g| g.is_some())
+        .count();
         if grammars > 1 {
             return Err(bad(
-                "Only one of regex, json_schema, or ebnf can be set".into()
+                "Only one of json_schema, regex, ebnf, or structural_tag can be set".into(),
             ));
         }
         // Not a Python restriction: the rust from_scheduler maps one rid to one response,
@@ -497,11 +545,6 @@ impl SamplingParams {
             )));
         }
         if let Some(beam_width) = self.beam_width {
-            if beam_width < 1 {
-                return Err(bad(format!(
-                    "beam_width must be at least 1, got {beam_width}."
-                )));
-            }
             // Also not a Python restriction: beam search returns its candidates
             // in `meta_info.beam_results`, which from_scheduler does not carry.
             if beam_width > 1 {
@@ -912,6 +955,43 @@ mod tests {
         assert_eq!(once.top_k, twice.top_k);
     }
 
+    #[test]
+    fn empty_grammar_constraints_are_unset() {
+        let sp = norm(r#"{"json_schema":"","regex":"","ebnf":"","structural_tag":""}"#);
+        assert!(sp.json_schema.is_none());
+        assert!(sp.regex.is_none());
+        assert!(sp.ebnf.is_none());
+        assert!(sp.structural_tag.is_none());
+    }
+
+    #[test]
+    fn structural_tag_is_mutually_exclusive_with_other_grammars() {
+        for field in ["json_schema", "regex", "ebnf"] {
+            let json = format!(r#"{{"structural_tag":"tag","{field}":"other"}}"#);
+            assert!(norm_err(&json).to_string().contains("Only one of"));
+        }
+    }
+
+    #[test]
+    fn custom_params_matches_python_shape() {
+        assert!(
+            norm(r#"{"custom_params":{"null":null,"bool":true,"int":1,"float":1.5,"str":"x","list":[1,"x",null],"object":{"x":1}}}"#)
+                .custom_params
+                .is_some()
+        );
+
+        for json in [
+            r#"{"custom_params":[]}"#,
+            r#"{"custom_params":{"nested_list":[[1]]}}"#,
+            r#"{"custom_params":{"nested_object":{"x":{"y":1}}}}"#,
+        ] {
+            assert!(
+                serde_json::from_str::<SamplingParams>(json).is_err(),
+                "{json} must not parse"
+            );
+        }
+    }
+
     /// `skip_tokenizer_init` has no tokenizer, so the text-matching stop features
     /// and `min_new_tokens` (needs eos_token_id) are 400s, not silent no-ops.
     /// Mirrors Python `raise_if_tokenizer_required`.