sglang rust server egress message (#32342)

rust/sglang-server/src/message.rs

@@ -8,6 +8,8 @@
 //! `SamplingParams` port), [`io_struct`] (the scheduler wire structs), [`types`]
 //! (the shared wire-shape adapters both directions use).
 
+mod egress;
+mod finish_reason;
 mod io_struct;
 mod request;
 mod sampling;

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

@@ -0,0 +1,217 @@
+//! The terminal finish reason: Python's `FinishReasonDict` — what
+//! `BaseFinishReason.to_json()` (schedule_batch.py) puts on the egress wire, and
+//! what the API echoes back as `meta_info.finish_reason`. Ingress has no
+//! counterpart; it rides in the [`BatchHeader`](super::egress::BatchHeader) and on
+//! each terminal [`ChunkEvent`](super::ChunkEvent).
+
+use serde::{Deserialize, Serialize};
+
+/// The stop that ended a request — Python's `matched` key, typed
+/// `Union[str, int, List[int]]`: a stop token id, a stop string
+/// (`FINISH_MATCHED_STR` / `FINISHED_MATCHED_REGEX`), or a multi-token stop
+/// sequence. Untagged because the three wire shapes are disjoint, so the shape
+/// alone picks the arm.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum Matched {
+    Token(i64),
+    Str(String),
+    Tokens(Vec<i64>),
+}
+
+/// The finish reasons this build knows, keyed by the `type` tag Python writes.
+/// Every field is optional so a reason missing one still classifies (and keeps its
+/// tag) instead of falling through to [`FinishReason::Unknown`].
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "lowercase")]
+pub enum FinishKind {
+    /// `FINISH_MATCHED_TOKEN` / `FINISH_MATCHED_STR` / `FINISHED_MATCHED_REGEX` —
+    /// all three report `type: "stop"`, to match the OpenAI API's value.
+    Stop {
+        #[serde(default)]
+        matched: Option<Matched>,
+    },
+    /// `FINISH_LENGTH` — hit `max_new_tokens` (or the context limit).
+    Length {
+        #[serde(default)]
+        length: Option<u64>,
+    },
+    /// `FINISH_ABORT` — a scheduler-side termination. Boxed because it is by far
+    /// the widest variant and the rarest: unboxed it would set the size of every
+    /// [`ChunkEvent`], including the plain stop/length ones (see
+    /// `chunk_event_frame_stays_small`).
+    Abort(Box<AbortReason>),
+}
+
+/// The `FINISH_ABORT` payload. `status_code`/`err_type` are `None` for a plain
+/// abort and set for a request error (e.g. over-context → 400); Python emits all
+/// three keys, nulls included, so none of them is skipped on the way out.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct AbortReason {
+    #[serde(default)]
+    pub message: Option<String>,
+    #[serde(default)]
+    pub status_code: Option<u16>,
+    #[serde(default)]
+    pub err_type: Option<String>,
+}
+
+/// A terminal finish reason exactly as `BaseFinishReason.to_json()`
+/// (schedule_batch.py) puts it on the wire — Python's `FinishReasonDict`. It
+/// round-trips: the API echoes it verbatim as `meta_info.finish_reason`, so the
+/// serialized form must stay key-for-key what Python would have sent.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum FinishReason {
+    Known(FinishKind),
+    /// A `type` this build doesn't know, kept as its raw map and echoed unchanged.
+    /// This arm is why the outer enum is untagged: a finish reason added Python-side
+    /// must not fail the header decode, which rejects the whole frame — every
+    /// request in the batch, not just the one that carried it.
+    Unknown(serde_json::Map<String, serde_json::Value>),
+}
+
+impl From<FinishKind> for FinishReason {
+    fn from(kind: FinishKind) -> Self {
+        FinishReason::Known(kind)
+    }
+}
+
+impl FinishReason {
+    /// The stop this request matched, if it stopped on one. `None` for
+    /// length/abort and for an unknown type.
+    pub fn matched(&self) -> Option<&Matched> {
+        match self {
+            FinishReason::Known(FinishKind::Stop { matched }) => matched.as_ref(),
+            _ => None,
+        }
+    }
+
+    /// `Some((status, message))` when this is an abort carrying a `status_code` —
+    /// a scheduler-side request error the API surfaces as that HTTP status instead
+    /// of as a normal completion. A plain abort (no code) reads as `None`.
+    pub fn abort_status(&self) -> Option<(u16, &str)> {
+        match self {
+            FinishReason::Known(FinishKind::Abort(a)) => Some((
+                a.status_code?,
+                a.message.as_deref().unwrap_or("request aborted"),
+            )),
+            _ => None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Parse a finish reason from the map Python puts on the wire, so these tests
+    /// exercise the real deserialization and not just the classifier.
+    fn fr(v: serde_json::Value) -> Option<FinishReason> {
+        Some(serde_json::from_value(v).expect("finish reason must parse"))
+    }
+
+    /// Classify through an `Option`, the shape both API paths actually hold.
+    /// Owned message so a caller can classify a temporary inline.
+    fn abort_status(f: &Option<FinishReason>) -> Option<(u16, String)> {
+        let (code, message) = f.as_ref()?.abort_status()?;
+        Some((code, message.to_string()))
+    }
+
+    /// The classifier both API paths use: a validation abort yields its
+    /// `(code, message)` (the streaming path turns this into an SSE error event
+    /// instead of a normal `Done` frame); anything else yields `None`.
+    #[test]
+    fn abort_status_extracts_code_and_message() {
+        let (code, msg) = abort_status(&fr(serde_json::json!({
+            "type": "abort", "message": "over the limit", "status_code": 400
+        })))
+        .expect("validation abort → (code, message)");
+        assert_eq!(code, 400);
+        assert_eq!(msg, "over the limit");
+        // Normal finish, bare abort (no status), and no finish → not an error.
+        assert!(abort_status(&fr(serde_json::json!({"type": "stop"}))).is_none());
+        assert!(abort_status(&fr(serde_json::json!({"type": "abort"}))).is_none());
+        assert!(abort_status(&None).is_none());
+    }
+
+    /// A normal finish, a bare abort (no status), and no finish are not errors
+    /// (the unary path returns them as a 200 result frame).
+    #[test]
+    fn non_error_finishes_stay_ok() {
+        assert!(abort_status(&fr(serde_json::json!({"type": "stop", "matched": 5}))).is_none());
+        assert!(abort_status(&fr(serde_json::json!({"type": "length", "length": 8}))).is_none());
+        assert!(
+            abort_status(&fr(
+                serde_json::json!({"type": "abort", "message": "Aborted"})
+            ))
+            .is_none()
+        );
+        assert!(abort_status(&None).is_none());
+    }
+
+    /// The `matched` accessor reads only a stop, and reads every shape of one.
+    #[test]
+    fn matched_reads_stops_only() {
+        let m = |v| fr(v).unwrap().matched().cloned();
+        assert_eq!(
+            m(serde_json::json!({"type": "stop", "matched": 9})),
+            Some(Matched::Token(9))
+        );
+        assert_eq!(
+            m(serde_json::json!({"type": "stop", "matched": "</s>"})),
+            Some(Matched::Str("</s>".into()))
+        );
+        assert_eq!(
+            m(serde_json::json!({"type": "stop", "matched": [9, 10]})),
+            Some(Matched::Tokens(vec![9, 10]))
+        );
+        // A stop with no `matched`, and the non-stop reasons.
+        assert_eq!(m(serde_json::json!({"type": "stop"})), None);
+        assert_eq!(m(serde_json::json!({"type": "length", "length": 8})), None);
+        assert_eq!(m(serde_json::json!({"type": "abort"})), None);
+    }
+
+    /// Python's abort dict always carries `status_code`/`err_type`, null included,
+    /// and the reason is echoed verbatim into `meta_info.finish_reason` — so the
+    /// typed form must serialize back key-for-key, nulls and all. A dropped null
+    /// key is a silently changed response body.
+    #[test]
+    fn finish_reason_round_trips_python_shapes() {
+        for wire in [
+            serde_json::json!({"type": "stop", "matched": 9}),
+            serde_json::json!({"type": "stop", "matched": "</s>"}),
+            serde_json::json!({"type": "stop", "matched": [9, 10]}),
+            serde_json::json!({"type": "length", "length": 8}),
+            serde_json::json!({
+                "type": "abort", "message": "Aborted",
+                "status_code": null, "err_type": null
+            }),
+            serde_json::json!({
+                "type": "abort", "message": "over the limit",
+                "status_code": 400, "err_type": "BadRequestError"
+            }),
+        ] {
+            let parsed: FinishReason = serde_json::from_value(wire.clone()).unwrap();
+            assert!(
+                matches!(parsed, FinishReason::Known(_)),
+                "must classify, not fall back to Unknown: {wire}"
+            );
+            assert_eq!(serde_json::to_value(&parsed).unwrap(), wire);
+        }
+    }
+
+    /// A `type` this build doesn't know must not fail the decode — the header holds
+    /// the whole batch, so a rejected frame drops every request in it, not just the
+    /// one carrying the new reason. It is kept verbatim and echoed unchanged.
+    #[test]
+    fn unknown_finish_type_is_preserved_not_rejected() {
+        let wire = serde_json::json!({"type": "tool_calls", "name": "search"});
+        let parsed: FinishReason = serde_json::from_value(wire.clone()).unwrap();
+        assert!(matches!(parsed, FinishReason::Unknown(_)));
+        assert_eq!(serde_json::to_value(&parsed).unwrap(), wire);
+        // Unknown reasons classify as neither an abort nor a matched stop.
+        assert!(parsed.abort_status().is_none());
+        assert!(parsed.matched().is_none());
+    }
+}