fix(protocol): forward sampling masks for chat completions

src/protocols/chat.rs

@@ -360,6 +360,10 @@ pub struct ChatCompletionRequest {
     #[serde(default)]
     pub return_routed_experts: bool,
 
+    /// Return token sampling masks in response metadata (SGLang extension)
+    #[serde(default)]
+    pub return_sampling_mask: bool,
+
     /// Absolute start position for returned routed experts (SGLang extension)
     #[serde(default)]
     pub routed_experts_start_len: i32,

tests/routing/test_openai_routing.rs

@@ -618,6 +618,7 @@ async fn test_unsupported_endpoints() {
         return_hidden_states: false,
         return_routed_experts: false,
         return_sampling_mask: None,
+        return_indexer_topk: false,
         routed_experts_start_len: 0,
         return_prompt_token_ids: false,
         require_reasoning: false,
@@ -696,6 +697,61 @@ async fn test_openai_router_chat_completion_with_mock() {
     assert!(!chat_response["choices"].as_array().unwrap().is_empty());
 }
 
+#[tokio::test]
+async fn test_openai_router_forwards_return_sampling_mask() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let received = Arc::new(tokio::sync::Mutex::new(None));
+    let app = Router::new().route(
+        "/v1/chat/completions",
+        post({
+            let received = received.clone();
+            move |Json(request): Json<serde_json::Value>| {
+                let received = received.clone();
+                async move {
+                    *received.lock().await = Some(request);
+                    Json(json!({
+                        "id": "chatcmpl-sampling-mask",
+                        "object": "chat.completion",
+                        "created": 1_700_000_000,
+                        "model": "test-model",
+                        "choices": [{
+                            "index": 0,
+                            "message": {"role": "assistant", "content": "ok"},
+                            "finish_reason": "stop"
+                        }]
+                    }))
+                }
+            }
+        }),
+    );
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    let ctx = crate::common::test_app::create_test_app_context().await;
+    crate::common::test_app::register_external_worker(
+        &ctx,
+        &format!("http://{addr}"),
+        Some(vec!["test-model"]),
+    );
+    let router = OpenAIRouter::new(&ctx).await.unwrap();
+    let mut request = create_minimal_chat_request();
+    request.model = "test-model".to_string();
+    request.return_sampling_mask = true;
+
+    let response = router.route_chat(None, &request, None).await;
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let forwarded = received
+        .lock()
+        .await
+        .take()
+        .expect("worker received request");
+    assert_eq!(forwarded["return_sampling_mask"], true);
+    server.abort();
+}
+
 /// Test full E2E flow with Axum server
 #[tokio::test]
 async fn test_openai_e2e_with_server() {

tests/spec/chat_completion.rs

@@ -578,6 +578,7 @@ fn test_sglang_extension_fields_roundtrip() {
         "messages": [{"role": "user", "content": "hello"}],
         "return_hidden_states": true,
         "return_routed_experts": true,
+        "return_sampling_mask": true,
         "routed_experts_start_len": 10,
         "return_cached_tokens_details": true,
         "return_prompt_token_ids": true,
@@ -604,6 +605,7 @@ fn test_sglang_extension_fields_roundtrip() {
         serde_json::from_value(json_with_extensions).expect("should deserialize");
     assert!(req.return_hidden_states);
     assert!(req.return_routed_experts);
+    assert!(req.return_sampling_mask);
     assert_eq!(req.routed_experts_start_len, 10);
     assert!(req.return_cached_tokens_details);
     assert!(req.return_prompt_token_ids);
@@ -619,6 +621,7 @@ fn test_sglang_extension_fields_roundtrip() {
     let serialized = serde_json::to_value(&req).expect("should serialize");
     assert_eq!(serialized["return_hidden_states"], true);
     assert_eq!(serialized["return_routed_experts"], true);
+    assert_eq!(serialized["return_sampling_mask"], true);
     assert_eq!(serialized["routed_experts_start_len"], 10);
     assert_eq!(serialized["return_cached_tokens_details"], true);
     assert_eq!(serialized["return_prompt_token_ids"], true);
@@ -652,6 +655,7 @@ fn test_sglang_extension_fields_default_values() {
         serde_json::from_value(json_minimal).expect("should deserialize");
     assert!(!req.return_hidden_states);
     assert!(!req.return_routed_experts);
+    assert!(!req.return_sampling_mask);
     assert_eq!(req.routed_experts_start_len, 0);
     assert!(!req.return_cached_tokens_details);
     assert!(!req.return_prompt_token_ids);