[Router]fix: fix get_load missing api_key (#10385)

sgl-router/src/core/worker_builder.rs

@@ -11,6 +11,7 @@ pub struct BasicWorkerBuilder {
     url: String,
 
     // Optional fields with defaults
+    api_key: Option<String>,
     worker_type: WorkerType,
     connection_mode: ConnectionMode,
     labels: HashMap<String, String>,
@@ -24,6 +25,7 @@ impl BasicWorkerBuilder {
     pub fn new(url: impl Into<String>) -> Self {
         Self {
             url: url.into(),
+            api_key: None,
             worker_type: WorkerType::Regular,
             connection_mode: ConnectionMode::Http,
             labels: HashMap::new(),
@@ -37,6 +39,7 @@ impl BasicWorkerBuilder {
     pub fn new_with_type(url: impl Into<String>, worker_type: WorkerType) -> Self {
         Self {
             url: url.into(),
+            api_key: None,
             worker_type,
             connection_mode: ConnectionMode::Http,
             labels: HashMap::new(),
@@ -46,6 +49,12 @@ impl BasicWorkerBuilder {
         }
     }
 
+    /// Set the API key
+    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
+        self.api_key = Some(api_key.into());
+        self
+    }
+
     /// Set the worker type (Regular, Prefill, or Decode)
     pub fn worker_type(mut self, worker_type: WorkerType) -> Self {
         self.worker_type = worker_type;
@@ -98,6 +107,7 @@ impl BasicWorkerBuilder {
 
         let metadata = WorkerMetadata {
             url: self.url.clone(),
+            api_key: self.api_key,
             worker_type: self.worker_type,
             connection_mode: self.connection_mode,
             labels: self.labels,
@@ -121,6 +131,7 @@ impl BasicWorkerBuilder {
 pub struct DPAwareWorkerBuilder {
     // Required fields
     base_url: String,
+    api_key: Option<String>,
     dp_rank: usize,
     dp_size: usize,
 
@@ -138,6 +149,7 @@ impl DPAwareWorkerBuilder {
     pub fn new(base_url: impl Into<String>, dp_rank: usize, dp_size: usize) -> Self {
         Self {
             base_url: base_url.into(),
+            api_key: None,
             dp_rank,
             dp_size,
             worker_type: WorkerType::Regular,
@@ -158,6 +170,7 @@ impl DPAwareWorkerBuilder {
     ) -> Self {
         Self {
             base_url: base_url.into(),
+            api_key: None,
             dp_rank,
             dp_size,
             worker_type,
@@ -169,6 +182,12 @@ impl DPAwareWorkerBuilder {
         }
     }
 
+    /// Set the API key
+    pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
+        self.api_key = Some(api_key.into());
+        self
+    }
+
     /// Set the worker type (Regular, Prefill, or Decode)
     pub fn worker_type(mut self, worker_type: WorkerType) -> Self {
         self.worker_type = worker_type;
@@ -228,6 +247,10 @@ impl DPAwareWorkerBuilder {
         if let Some(client) = self.grpc_client {
             builder = builder.grpc_client(client);
         }
+        // Add API key if provided
+        if let Some(api_key) = self.api_key {
+            builder = builder.api_key(api_key);
+        }
 
         let base_worker = builder.build();
 
@@ -382,6 +405,7 @@ mod tests {
             .connection_mode(ConnectionMode::Http)
             .labels(labels.clone())
             .health_config(health_config.clone())
+            .api_key("test_api_key")
             .build();
 
         assert_eq!(worker.url(), "http://localhost:8080@3");