From cf6800f6cce01c567eded824b6ef1a10d45b0bfb Mon Sep 17 00:00:00 2001
From: Simo Lin <linsimo.mark@gmail.com>
Date: Sun, 4 Jan 2026 10:07:55 -0800
Subject: [PATCH] [model-gateway] optimize Vec and HashMap allocations in
 responses api (#16406)

---
 sgl-model-gateway/src/routers/grpc/utils.rs | 24 +++++++++++----------
 sgl-model-gateway/src/routers/openai/mcp.rs |  5 +++--
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/sgl-model-gateway/src/routers/grpc/utils.rs b/sgl-model-gateway/src/routers/grpc/utils.rs
index 83e3d11b4..29b777f02 100644
--- a/sgl-model-gateway/src/routers/grpc/utils.rs
+++ b/sgl-model-gateway/src/routers/grpc/utils.rs
@@ -279,8 +279,7 @@ pub fn generate_tool_constraints(
 /// Build JSON schema for required tool calls (array with minItems: 1)
 /// Includes $defs consolidation from all tools (matching Python's behavior)
 fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
-    // Build anyOf schemas for each tool
-    let mut any_of_schemas = Vec::new();
+    let mut any_of_schemas = Vec::with_capacity(tools.len());
     for tool in tools {
         let tool_schema = json!({
             "properties": {
@@ -296,7 +295,7 @@ fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
     }
 
     // Consolidate $defs from all tools (matching Python's _get_tool_schema_defs)
-    let mut all_defs: HashMap<String, Value> = HashMap::new();
+    let mut all_defs: Map<String, Value> = Map::new();
     for tool in tools {
         if let Value::Object(params) = &tool.function.parameters {
             if let Some(Value::Object(defs)) = params.get("$defs") {
@@ -332,8 +331,7 @@ fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
     // Add $defs if any were found (matching Python's behavior)
     if !all_defs.is_empty() {
         if let Value::Object(ref mut schema_obj) = array_schema {
-            let defs_value = Value::Object(all_defs.into_iter().collect::<Map<String, Value>>());
-            schema_obj.insert("$defs".to_string(), defs_value);
+            schema_obj.insert("$defs".to_string(), Value::Object(all_defs));
         }
     }
 
@@ -439,8 +437,8 @@ pub fn process_chat_messages(
             .transpose()
             .map_err(|e| format!("Failed to serialize tools: {}", e))?;
 
-        // Build template kwargs, merging reasoning_effort if present
-        let mut combined_template_kwargs = HashMap::new();
+        let kwargs_capacity = 1 + request.chat_template_kwargs.as_ref().map_or(0, |k| k.len());
+        let mut combined_template_kwargs = HashMap::with_capacity(kwargs_capacity);
 
         // Add reasoning_effort if present (like Python does)
         if let Some(reasoning_effort) = &request.reasoning_effort {
@@ -878,7 +876,7 @@ pub fn convert_proto_to_openai_logprobs(
     proto_logprobs: &OutputLogProbs,
     tokenizer: &Arc<dyn Tokenizer>,
 ) -> Result<ChatLogProbs, String> {
-    let mut content_items = Vec::new();
+    let mut content_items = Vec::with_capacity(proto_logprobs.token_logprobs.len());
 
     // Decode token IDs to text (always with skip_special_tokens=false for logprobs)
     let token_texts: Vec<String> = proto_logprobs
@@ -901,8 +899,9 @@ pub fn convert_proto_to_openai_logprobs(
         let bytes = Some(token_text.as_bytes().to_vec());
 
         // Build top_logprobs for this position
-        let mut top_logprobs = Vec::new();
-        if let Some(top_logprobs_entry) = proto_logprobs.top_logprobs.get(i) {
+        let top_logprobs = if let Some(top_logprobs_entry) = proto_logprobs.top_logprobs.get(i) {
+            let mut top_logprobs = Vec::with_capacity(top_logprobs_entry.values.len());
+
             // Decode top token IDs (always with skip_special_tokens=false)
             let top_token_texts: Vec<String> = top_logprobs_entry
                 .token_ids
@@ -928,7 +927,10 @@ pub fn convert_proto_to_openai_logprobs(
                     });
                 }
             }
-        }
+            top_logprobs
+        } else {
+            Vec::new()
+        };
 
         content_items.push(ChatLogProbsContent {
             token: token_text,
diff --git a/sgl-model-gateway/src/routers/openai/mcp.rs b/sgl-model-gateway/src/routers/openai/mcp.rs
index 47a2b1dab..95d07597f 100644
--- a/sgl-model-gateway/src/routers/openai/mcp.rs
+++ b/sgl-model-gateway/src/routers/openai/mcp.rs
@@ -217,8 +217,8 @@ pub(super) fn prepare_mcp_payload_for_streaming(
         }
 
         // Build function tools for all discovered MCP tools
-        let mut tools_json = Vec::new();
         let tools = active_mcp.list_tools();
+        let mut tools_json = Vec::with_capacity(tools.len());
         for t in tools {
             let parameters = Value::Object((*t.input_schema).clone());
             let tool = serde_json::json!({
@@ -252,7 +252,8 @@ pub(super) fn build_resume_payload(
         .ok_or_else(|| "payload not an object".to_string())?;
 
     // Build input array: start with original user input
-    let mut input_array = Vec::new();
+    // Pre-allocate: 1 for user message + conversation history
+    let mut input_array = Vec::with_capacity(1 + conversation_history.len());
 
     // Add original user message
     // For structured input, serialize the original input items