diff --git a/sgl-model-gateway/src/routers/openai/responses/mcp.rs b/sgl-model-gateway/src/routers/openai/responses/mcp.rs index 8f2ae7a62..75425ba90 100644 --- a/sgl-model-gateway/src/routers/openai/responses/mcp.rs +++ b/sgl-model-gateway/src/routers/openai/responses/mcp.rs @@ -198,8 +198,8 @@ pub(super) async fn execute_streaming_tool_calls( // Payload Transformation // ============================================================================ -/// Transform payload to replace MCP tools with function tools for streaming -pub(super) fn prepare_mcp_payload_for_streaming( +/// Transform payload to replace MCP tools with function tools +pub(super) fn prepare_mcp_tools_as_functions( payload: &mut Value, active_mcp: &Arc, server_keys: &[String], diff --git a/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs b/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs index f02d9f3e2..51c7ae3fd 100644 --- a/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs @@ -11,8 +11,8 @@ use serde_json::{json, Value}; use tracing::warn; use super::{ - mcp::{execute_tool_loop, prepare_mcp_payload_for_streaming}, - utils::{mask_tools_as_mcp, patch_streaming_response_json}, + mcp::{execute_tool_loop, prepare_mcp_tools_as_functions}, + utils::{mask_tools_as_mcp, patch_response_with_request_metadata}, }; use crate::routers::{ header_utils::{apply_provider_headers, extract_auth_header}, @@ -71,7 +71,7 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response server_keys: server_keys.clone(), ..McpLoopConfig::default() }; - prepare_mcp_payload_for_streaming(&mut payload, mcp, &server_keys); + prepare_mcp_tools_as_functions(&mut payload, mcp, &server_keys); match execute_tool_loop( ctx.components.client(), @@ -140,7 +140,7 @@ pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response } mask_tools_as_mcp(&mut response_json, original_body); - patch_streaming_response_json( + patch_response_with_request_metadata( &mut response_json, original_body, previous_response_id.as_deref(), diff --git a/sgl-model-gateway/src/routers/openai/responses/streaming.rs b/sgl-model-gateway/src/routers/openai/responses/streaming.rs index 2176d5cf4..2c820517b 100644 --- a/sgl-model-gateway/src/routers/openai/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/streaming.rs @@ -26,10 +26,10 @@ use super::{ common::{extract_output_index, get_event_type, parse_sse_block, ChunkProcessor}, mcp::{ build_resume_payload, execute_streaming_tool_calls, inject_mcp_metadata_streaming, - prepare_mcp_payload_for_streaming, send_mcp_list_tools_events, ToolLoopState, + prepare_mcp_tools_as_functions, send_mcp_list_tools_events, ToolLoopState, }, tool_handler::{StreamAction, StreamingToolHandler}, - utils::{mask_tools_as_mcp, patch_streaming_response_json, rewrite_streaming_block}, + utils::{mask_tools_as_mcp, patch_response_with_request_metadata, rewrite_streaming_block}, }; use crate::{ protocols::{ @@ -452,7 +452,7 @@ pub(super) fn send_final_response_event( } mask_tools_as_mcp(&mut final_response, ctx.original_request); - patch_streaming_response_json( + patch_response_with_request_metadata( &mut final_response, ctx.original_request, ctx.previous_response_id, @@ -590,7 +590,7 @@ pub(super) async fn handle_simple_streaming_passthrough( } let encountered_error = accumulator.encountered_error().cloned(); if let Some(mut response_json) = accumulator.into_final_response() { - patch_streaming_response_json( + patch_response_with_request_metadata( &mut response_json, &original_request, previous_response_id.as_deref(), @@ -642,7 +642,7 @@ pub(super) async fn handle_streaming_with_tool_interception( ) -> Response { // Transform MCP tools to function tools in payload let mut payload = req.payload; - prepare_mcp_payload_for_streaming(&mut payload, active_mcp, &server_keys); + prepare_mcp_tools_as_functions(&mut payload, active_mcp, &server_keys); let (tx, rx) = mpsc::unbounded_channel::>(); let should_store = req.original_body.store.unwrap_or(false); @@ -885,7 +885,7 @@ pub(super) async fn handle_streaming_with_tool_interception( ); mask_tools_as_mcp(&mut response_json, &original_request); - patch_streaming_response_json( + patch_response_with_request_metadata( &mut response_json, &original_request, previous_response_id.as_deref(), diff --git a/sgl-model-gateway/src/routers/openai/responses/utils.rs b/sgl-model-gateway/src/routers/openai/responses/utils.rs index f3b51ef5b..d85bdfad2 100644 --- a/sgl-model-gateway/src/routers/openai/responses/utils.rs +++ b/sgl-model-gateway/src/routers/openai/responses/utils.rs @@ -26,8 +26,17 @@ where } } -/// Patch streaming response JSON with metadata from original request -pub(super) fn patch_streaming_response_json( +/// Patch response JSON with metadata from original request +/// +/// The upstream response may be missing fields that were in the original request. +/// This function ensures these fields are preserved in the final response: +/// - `previous_response_id` - conversation threading +/// - `instructions` - system instructions +/// - `metadata` - user-provided metadata +/// - `store` - whether to persist the response +/// - `model` - model identifier +/// - `safety_identifier` - user identifier for safety +pub(super) fn patch_response_with_request_metadata( response_json: &mut Value, original_body: &ResponsesRequest, original_previous_response_id: Option<&str>,