[router][openai] Rename prepare_mcp_payload_for_streaming and patch_streaming_response_json (#16596)

This commit is contained in:
Chang Su
2026-01-06 16:33:29 -08:00
committed by GitHub
parent 05b54b6d7b
commit fb5b71d015
4 changed files with 23 additions and 14 deletions
@@ -198,8 +198,8 @@ pub(super) async fn execute_streaming_tool_calls(
// Payload Transformation // Payload Transformation
// ============================================================================ // ============================================================================
/// Transform payload to replace MCP tools with function tools for streaming /// Transform payload to replace MCP tools with function tools
pub(super) fn prepare_mcp_payload_for_streaming( pub(super) fn prepare_mcp_tools_as_functions(
payload: &mut Value, payload: &mut Value,
active_mcp: &Arc<mcp::McpManager>, active_mcp: &Arc<mcp::McpManager>,
server_keys: &[String], server_keys: &[String],
@@ -11,8 +11,8 @@ use serde_json::{json, Value};
use tracing::warn; use tracing::warn;
use super::{ use super::{
mcp::{execute_tool_loop, prepare_mcp_payload_for_streaming}, mcp::{execute_tool_loop, prepare_mcp_tools_as_functions},
utils::{mask_tools_as_mcp, patch_streaming_response_json}, utils::{mask_tools_as_mcp, patch_response_with_request_metadata},
}; };
use crate::routers::{ use crate::routers::{
header_utils::{apply_provider_headers, extract_auth_header}, 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(), server_keys: server_keys.clone(),
..McpLoopConfig::default() ..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( match execute_tool_loop(
ctx.components.client(), 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); mask_tools_as_mcp(&mut response_json, original_body);
patch_streaming_response_json( patch_response_with_request_metadata(
&mut response_json, &mut response_json,
original_body, original_body,
previous_response_id.as_deref(), previous_response_id.as_deref(),
@@ -26,10 +26,10 @@ use super::{
common::{extract_output_index, get_event_type, parse_sse_block, ChunkProcessor}, common::{extract_output_index, get_event_type, parse_sse_block, ChunkProcessor},
mcp::{ mcp::{
build_resume_payload, execute_streaming_tool_calls, inject_mcp_metadata_streaming, 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}, 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::{ use crate::{
protocols::{ protocols::{
@@ -452,7 +452,7 @@ pub(super) fn send_final_response_event(
} }
mask_tools_as_mcp(&mut final_response, ctx.original_request); mask_tools_as_mcp(&mut final_response, ctx.original_request);
patch_streaming_response_json( patch_response_with_request_metadata(
&mut final_response, &mut final_response,
ctx.original_request, ctx.original_request,
ctx.previous_response_id, ctx.previous_response_id,
@@ -590,7 +590,7 @@ pub(super) async fn handle_simple_streaming_passthrough(
} }
let encountered_error = accumulator.encountered_error().cloned(); let encountered_error = accumulator.encountered_error().cloned();
if let Some(mut response_json) = accumulator.into_final_response() { if let Some(mut response_json) = accumulator.into_final_response() {
patch_streaming_response_json( patch_response_with_request_metadata(
&mut response_json, &mut response_json,
&original_request, &original_request,
previous_response_id.as_deref(), previous_response_id.as_deref(),
@@ -642,7 +642,7 @@ pub(super) async fn handle_streaming_with_tool_interception(
) -> Response { ) -> Response {
// Transform MCP tools to function tools in payload // Transform MCP tools to function tools in payload
let mut payload = req.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::<Result<Bytes, io::Error>>(); let (tx, rx) = mpsc::unbounded_channel::<Result<Bytes, io::Error>>();
let should_store = req.original_body.store.unwrap_or(false); 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); mask_tools_as_mcp(&mut response_json, &original_request);
patch_streaming_response_json( patch_response_with_request_metadata(
&mut response_json, &mut response_json,
&original_request, &original_request,
previous_response_id.as_deref(), previous_response_id.as_deref(),
@@ -26,8 +26,17 @@ where
} }
} }
/// Patch streaming response JSON with metadata from original request /// Patch response JSON with metadata from original request
pub(super) fn patch_streaming_response_json( ///
/// 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, response_json: &mut Value,
original_body: &ResponsesRequest, original_body: &ResponsesRequest,
original_previous_response_id: Option<&str>, original_previous_response_id: Option<&str>,