From 8db2802b2da06f9d4c4ad8db2fb9ec3c06e5d156 Mon Sep 17 00:00:00 2001
From: Simo Lin <linsimo.mark@gmail.com>
Date: Sun, 25 Jan 2026 09:35:59 -0500
Subject: [PATCH] update to use official openai protocol crate (#17710)

---
 sgl-model-gateway/Cargo.toml                 |  1 +
 sgl-model-gateway/src/core/job_queue.rs      | 41 -----------------
 sgl-model-gateway/src/core/worker_manager.rs | 47 +-------------------
 sgl-model-gateway/src/lib.rs                 |  2 +-
 4 files changed, 4 insertions(+), 87 deletions(-)

diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml
index bc6e5787f..601bc5271 100644
--- a/sgl-model-gateway/Cargo.toml
+++ b/sgl-model-gateway/Cargo.toml
@@ -84,6 +84,7 @@ anyhow = "1.0"
 tokenizers = { version = "0.22.0" }
 tiktoken-rs = { version = "0.7.0" }
 reasoning-parser = "1.0.0"
+openai-protocol = { version = "1.0.0", features = ["axum"] }
 minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins"] }
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
 rustls = { version = "0.23", default-features = false, features = ["ring", "std"] }
diff --git a/sgl-model-gateway/src/core/job_queue.rs b/sgl-model-gateway/src/core/job_queue.rs
index bb928a836..e45b1dff6 100644
--- a/sgl-model-gateway/src/core/job_queue.rs
+++ b/sgl-model-gateway/src/core/job_queue.rs
@@ -99,47 +99,6 @@ impl Job {
     }
 }
 
-impl JobStatus {
-    fn pending(job_type: &str, worker_url: &str) -> Self {
-        Self {
-            job_type: job_type.to_string(),
-            worker_url: worker_url.to_string(),
-            status: "pending".to_string(),
-            message: None,
-            timestamp: SystemTime::now()
-                .duration_since(SystemTime::UNIX_EPOCH)
-                .unwrap()
-                .as_secs(),
-        }
-    }
-
-    fn processing(job_type: &str, worker_url: &str) -> Self {
-        Self {
-            job_type: job_type.to_string(),
-            worker_url: worker_url.to_string(),
-            status: "processing".to_string(),
-            message: None,
-            timestamp: SystemTime::now()
-                .duration_since(SystemTime::UNIX_EPOCH)
-                .unwrap()
-                .as_secs(),
-        }
-    }
-
-    fn failed(job_type: &str, worker_url: &str, error: String) -> Self {
-        Self {
-            job_type: job_type.to_string(),
-            worker_url: worker_url.to_string(),
-            status: "failed".to_string(),
-            message: Some(error),
-            timestamp: SystemTime::now()
-                .duration_since(SystemTime::UNIX_EPOCH)
-                .unwrap()
-                .as_secs(),
-        }
-    }
-}
-
 /// Job queue configuration
 #[derive(Clone, Debug)]
 pub struct JobQueueConfig {
diff --git a/sgl-model-gateway/src/core/worker_manager.rs b/sgl-model-gateway/src/core/worker_manager.rs
index 4845c7f80..4a21ae3fb 100644
--- a/sgl-model-gateway/src/core/worker_manager.rs
+++ b/sgl-model-gateway/src/core/worker_manager.rs
@@ -4,16 +4,13 @@
 
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
-use axum::{
-    response::{IntoResponse, Response},
-    Json,
-};
+use axum::response::{IntoResponse, Response};
 use futures::{
     future,
     stream::{self, StreamExt},
 };
 use http::StatusCode;
-use serde_json::{json, Value};
+use serde_json::Value;
 use tokio::{
     sync::{watch, Mutex},
     task::JoinHandle,
@@ -70,46 +67,6 @@ async fn fan_out(
         .await
 }
 
-impl IntoResponse for FlushCacheResult {
-    fn into_response(self) -> Response {
-        let status = if self.failed.is_empty() {
-            StatusCode::OK
-        } else {
-            StatusCode::PARTIAL_CONTENT
-        };
-
-        let mut body = json!({
-            "status": if self.failed.is_empty() { "success" } else { "partial_success" },
-            "message": self.message,
-            "workers_flushed": self.successful.len(),
-            "total_http_workers": self.http_workers,
-            "total_workers": self.total_workers
-        });
-
-        if !self.failed.is_empty() {
-            body["successful"] = json!(self.successful);
-            body["failed"] = json!(self
-                .failed
-                .into_iter()
-                .map(|(url, err)| json!({"worker": url, "error": err}))
-                .collect::<Vec<_>>());
-        }
-
-        (status, Json(body)).into_response()
-    }
-}
-
-impl IntoResponse for WorkerLoadsResult {
-    fn into_response(self) -> Response {
-        let loads: Vec<Value> = self
-            .loads
-            .iter()
-            .map(|info| json!({"worker": &info.worker, "load": info.load}))
-            .collect();
-        Json(json!({"workers": loads})).into_response()
-    }
-}
-
 pub enum EngineMetricsResult {
     Ok(String),
     Err(String),
diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs
index c3c28f909..20ff16cf0 100644
--- a/sgl-model-gateway/src/lib.rs
+++ b/sgl-model-gateway/src/lib.rs
@@ -10,7 +10,7 @@ pub mod middleware;
 pub mod multimodal;
 pub mod observability;
 pub mod policies;
-pub mod protocols;
+pub use openai_protocol as protocols;
 pub use reasoning_parser;
 pub mod routers;
 pub mod server;