Add manual routing policy for router (#15586)

sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs

@@ -61,15 +61,18 @@ impl PipelineStage for WorkerSelectionStage {
 
         // For Harmony, use selection_text produced during Harmony encoding
         // Otherwise, use original_text from regular preparation
-        let text = if prep.harmony_mode {
-            prep.selection_text.as_deref()
-        } else {
-            prep.original_text.as_deref()
+        let info = SelectWorkerInfo {
+            request_text: if prep.harmony_mode {
+                prep.selection_text.as_deref()
+            } else {
+                prep.original_text.as_deref()
+            },
+            routing_id: prep.routing_id.as_deref(),
         };
 
         let workers = match self.mode {
             WorkerSelectionMode::Regular => {
-                match self.select_single_worker(ctx.input.model_id.as_deref(), text) {
+                match self.select_single_worker(ctx.input.model_id.as_deref(), &info) {
                     Some(w) => WorkerSelection::Single { worker: w },
                     None => {
                         error!(
@@ -86,7 +89,7 @@ impl PipelineStage for WorkerSelectionStage {
                 }
             }
             WorkerSelectionMode::PrefillDecode => {
-                match self.select_pd_pair(ctx.input.model_id.as_deref(), text) {
+                match self.select_pd_pair(ctx.input.model_id.as_deref(), &info) {
                     Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode },
                     None => {
                         error!(
@@ -120,7 +123,7 @@ impl WorkerSelectionStage {
     fn select_single_worker(
         &self,
         model_id: Option<&str>,
-        text: Option<&str>,
+        info: &SelectWorkerInfo,
     ) -> Option<Arc<dyn Worker>> {
         // Get workers for the specified model, filtered by connection mode
         let workers = self.worker_registry.get_workers_filtered(
@@ -146,7 +149,7 @@ impl WorkerSelectionStage {
         };
 
         // Select worker using the policy
-        let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?;
+        let idx = policy.select_worker(&available, info)?;
         let selected = available[idx].clone();
 
         // Record worker selection metric
@@ -163,7 +166,7 @@ impl WorkerSelectionStage {
     fn select_pd_pair(
         &self,
         model_id: Option<&str>,
-        text: Option<&str>,
+        info: &SelectWorkerInfo,
     ) -> Option<(Arc<dyn Worker>, Arc<dyn Worker>)> {
         let all_workers = self.worker_registry.get_workers_filtered(
             model_id,
@@ -203,9 +206,8 @@ impl WorkerSelectionStage {
             None => self.policy_registry.get_default_policy(),
         };
 
-        let info = SelectWorkerInfo { request_text: text };
-        let prefill_idx = policy.select_worker(&available_prefill, &info)?;
-        let decode_idx = policy.select_worker(&available_decode, &info)?;
+        let prefill_idx = policy.select_worker(&available_prefill, info)?;
+        let decode_idx = policy.select_worker(&available_decode, info)?;
 
         let model = model_id.unwrap_or("default");
         let policy_name = policy.name();

sgl-model-gateway/src/routers/grpc/context.rs

@@ -94,6 +94,9 @@ pub struct PreparationOutput {
     /// Original text (for chat) or resolved text (for generate)
     pub original_text: Option<String>,
 
+    /// Routing ID for manual routing policy
+    pub routing_id: Option<String>,
+
     /// Tokenized input
     pub token_ids: Vec<u32>,
 

sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs

@@ -123,6 +123,7 @@ impl HarmonyPreparationStage {
         // Step 4: Store results
         ctx.state.preparation = Some(PreparationOutput {
             original_text: None,
+            routing_id: request.routing_id.clone(),
             token_ids: build_output.input_ids,
             processed_messages: None,
             tool_constraints,
@@ -203,6 +204,7 @@ impl HarmonyPreparationStage {
         // Step 4: Store results with constraint
         ctx.state.preparation = Some(PreparationOutput {
             original_text: None,
+            routing_id: request.routing_id.clone(),
             token_ids: build_output.input_ids,
             processed_messages: None,
             tool_constraints: constraint,

sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs

@@ -492,6 +492,7 @@ pub(super) async fn execute_tool_loop(
                 top_k: current_request.top_k,
                 min_p: current_request.min_p,
                 repetition_penalty: current_request.repetition_penalty,
+                routing_id: current_request.routing_id.clone(),
             };
 
             // Continue to next iteration
@@ -1070,6 +1071,7 @@ async fn execute_tool_loop_streaming_internal(
                 top_k: current_request.top_k,
                 min_p: current_request.min_p,
                 repetition_penalty: current_request.repetition_penalty,
+                routing_id: current_request.routing_id.clone(),
             };
 
             continue;

sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs

@@ -96,6 +96,7 @@ impl ChatPreparationStage {
         // Store results in context
         ctx.state.preparation = Some(PreparationOutput {
             original_text: Some(processed_messages.text.clone()),
+            routing_id: request.routing_id.clone(),
             token_ids,
             processed_messages: Some(processed_messages),
             tool_constraints: tool_call_constraint,

sgl-model-gateway/src/routers/grpc/regular/stages/embedding/preparation.rs

@@ -47,8 +47,9 @@ impl PipelineStage for EmbeddingPreparationStage {
             ));
         };
 
-        // Extract text from request
+        // Extract text and routing_id from request before borrowing ctx mutably
         let text = request.extract_text_for_routing();
+        let routing_id = request.routing_id.clone();
         if text.is_empty() {
             return Err(error::bad_request(
                 "empty_input",
@@ -77,6 +78,7 @@ impl PipelineStage for EmbeddingPreparationStage {
         // Store preparation output
         ctx.state.preparation = Some(PreparationOutput {
             original_text: Some(text),
+            routing_id,
             token_ids,
             processed_messages: None,
             tool_constraints: None,

sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs

@@ -68,6 +68,7 @@ impl GeneratePreparationStage {
 
         ctx.state.preparation = Some(PreparationOutput {
             original_text,
+            routing_id: request.routing_id.clone(),
             token_ids,
             processed_messages: None,
             tool_constraints: None,