From 5a2b1ed407c105213854a6b7d62ef0f1cabfba3f Mon Sep 17 00:00:00 2001 From: Chang Su Date: Mon, 5 Jan 2026 10:21:34 -0800 Subject: [PATCH] [grpc] Refactor grpc/regular/responses (#16509) --- .../routers/grpc/common/responses/handlers.rs | 2 +- .../routers/grpc/regular/responses/common.rs | 464 +++++++++ .../routers/grpc/regular/responses/context.rs | 19 +- .../grpc/regular/responses/handlers.rs | 736 +------------- .../src/routers/grpc/regular/responses/mod.rs | 27 +- .../grpc/regular/responses/non_streaming.rs | 417 ++++++++ .../responses/{tool_loop.rs => streaming.rs} | 921 +++++++----------- .../routers/grpc/regular/responses/types.rs | 18 - 8 files changed, 1328 insertions(+), 1276 deletions(-) create mode 100644 sgl-model-gateway/src/routers/grpc/regular/responses/common.rs create mode 100644 sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs rename sgl-model-gateway/src/routers/grpc/regular/responses/{tool_loop.rs => streaming.rs} (54%) delete mode 100644 sgl-model-gateway/src/routers/grpc/regular/responses/types.rs diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs b/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs index 0d9e84d3b..9b9f85bad 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs @@ -11,7 +11,7 @@ use tracing::{debug, error, warn}; use crate::{ data_connector::ResponseId, - routers::{error, grpc::regular::responses::context::ResponsesContext}, + routers::{error, grpc::regular::responses::ResponsesContext}, }; /// Implementation for GET /v1/responses/{response_id} diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs new file mode 100644 index 000000000..ecbc99c03 --- /dev/null +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs @@ -0,0 +1,464 @@ +//! Shared helpers and state tracking for Regular Responses +//! +//! This module contains common utilities used by both streaming and non-streaming paths: +//! - ToolLoopState for tracking multi-turn tool calling +//! - Helper functions for tool preparation and extraction +//! - MCP metadata builders +//! - Conversation history loading + +use std::sync::Arc; + +use axum::response::Response; +use serde_json::{json, Value}; +use tracing::{debug, warn}; +use uuid::Uuid; + +use super::context::ResponsesContext; +use crate::{ + data_connector::{self, ConversationId, ResponseId}, + mcp::{self, McpManager}, + protocols::{ + chat::ChatCompletionRequest, + common::{Function, Tool, ToolChoice, ToolChoiceValue}, + responses::{ + self, McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem, + ResponseOutputItem, ResponsesRequest, + }, + }, + routers::error, +}; + +// ============================================================================ +// Tool Loop State +// ============================================================================ + +/// State for tracking multi-turn tool calling loop +pub(super) struct ToolLoopState { + pub iteration: usize, + pub total_calls: usize, + pub conversation_history: Vec, + pub original_input: ResponseInput, + pub mcp_call_items: Vec, + pub server_label: String, +} + +impl ToolLoopState { + pub fn new(original_input: ResponseInput, server_label: String) -> Self { + Self { + iteration: 0, + total_calls: 0, + conversation_history: Vec::new(), + original_input, + mcp_call_items: Vec::new(), + server_label, + } + } + + pub fn record_call( + &mut self, + call_id: String, + tool_name: String, + args_json_str: String, + output_str: String, + success: bool, + error: Option, + ) { + // Add function_tool_call item with both arguments and output + self.conversation_history + .push(ResponseInputOutputItem::FunctionToolCall { + id: call_id.clone(), + call_id: call_id.clone(), + name: tool_name.clone(), + arguments: args_json_str.clone(), + output: Some(output_str.clone()), + status: Some("completed".to_string()), + }); + + // Add mcp_call output item for metadata + let mcp_call = build_mcp_call_item( + &tool_name, + &args_json_str, + &output_str, + &self.server_label, + success, + error.as_deref(), + ); + self.mcp_call_items.push(mcp_call); + } +} + +// ============================================================================ +// Tool Preparation and Extraction +// ============================================================================ + +/// Merge function tools from request with MCP tools and set tool_choice based on iteration +pub(super) fn prepare_chat_tools_and_choice( + chat_request: &mut ChatCompletionRequest, + mcp_chat_tools: &[Tool], + iteration: usize, +) { + // Merge function tools from request with MCP tools + let mut all_tools = chat_request.tools.clone().unwrap_or_default(); + all_tools.extend(mcp_chat_tools.iter().cloned()); + chat_request.tools = Some(all_tools); + + // Set tool_choice based on iteration + // - Iteration 0: Use user's tool_choice or default to auto + // - Iteration 1+: Always use auto to avoid infinite loops + chat_request.tool_choice = if iteration == 0 { + chat_request + .tool_choice + .clone() + .or(Some(ToolChoice::Value(ToolChoiceValue::Auto))) + } else { + Some(ToolChoice::Value(ToolChoiceValue::Auto)) + }; +} + +/// Extract all tool calls from chat response (for parallel tool call support) +pub(super) fn extract_all_tool_calls_from_chat( + response: &crate::protocols::chat::ChatCompletionResponse, +) -> Vec<(String, String, String)> { + // Check if response has choices with tool calls + let Some(choice) = response.choices.first() else { + return Vec::new(); + }; + let message = &choice.message; + + // Look for tool_calls in the message + if let Some(tool_calls) = &message.tool_calls { + tool_calls + .iter() + .map(|tool_call| { + ( + tool_call.id.clone(), + tool_call.function.name.clone(), + tool_call + .function + .arguments + .clone() + .unwrap_or_else(|| "{}".to_string()), + ) + }) + .collect() + } else { + Vec::new() + } +} + +/// Convert MCP tools to Chat API tool format +pub(super) fn convert_mcp_tools_to_chat_tools(mcp_tools: &[mcp::Tool]) -> Vec { + mcp_tools + .iter() + .map(|tool_info| Tool { + tool_type: "function".to_string(), + function: Function { + name: tool_info.name.to_string(), + description: tool_info.description.as_ref().map(|d| d.to_string()), + parameters: Value::Object((*tool_info.input_schema).clone()), + strict: None, + }, + }) + .collect() +} + +// ============================================================================ +// MCP Metadata Builders +// ============================================================================ + +/// Generate unique ID for MCP items +pub(super) fn generate_mcp_id(prefix: &str) -> String { + format!("{}_{}", prefix, Uuid::new_v4()) +} + +/// Build mcp_list_tools output item +pub(super) fn build_mcp_list_tools_item( + mcp: &Arc, + server_label: &str, +) -> ResponseOutputItem { + let tools = mcp.list_tools(); + let tools_info: Vec = tools + .iter() + .map(|t| McpToolInfo { + name: t.name.to_string(), + description: t.description.as_ref().map(|d| d.to_string()), + input_schema: Value::Object((*t.input_schema).clone()), + annotations: Some(json!({ + "read_only": false + })), + }) + .collect(); + + ResponseOutputItem::McpListTools { + id: generate_mcp_id("mcpl"), + server_label: server_label.to_string(), + tools: tools_info, + } +} + +/// Build mcp_call output item +pub(super) fn build_mcp_call_item( + tool_name: &str, + arguments: &str, + output: &str, + server_label: &str, + success: bool, + error: Option<&str>, +) -> ResponseOutputItem { + ResponseOutputItem::McpCall { + id: generate_mcp_id("mcp"), + status: if success { "completed" } else { "failed" }.to_string(), + approval_request_id: None, + arguments: arguments.to_string(), + error: error.map(|e| e.to_string()), + name: tool_name.to_string(), + output: output.to_string(), + server_label: server_label.to_string(), + } +} + +// ============================================================================ +// Conversation History Loading +// ============================================================================ + +/// Load conversation history and response chains, returning modified request +pub(super) async fn load_conversation_history( + ctx: &ResponsesContext, + request: &ResponsesRequest, +) -> Result { + let mut modified_request = request.clone(); + let mut conversation_items: Option> = None; + + // Handle previous_response_id by loading response chain + if let Some(ref prev_id_str) = modified_request.previous_response_id { + let prev_id = ResponseId::from(prev_id_str.as_str()); + match ctx + .response_storage + .get_response_chain(&prev_id, None) + .await + { + Ok(chain) => { + let mut items = Vec::new(); + for stored in chain.responses.iter() { + // Convert input items from stored input (which is now a JSON array) + if let Some(input_arr) = stored.input.as_array() { + for item in input_arr { + match serde_json::from_value::(item.clone()) { + Ok(input_item) => { + items.push(input_item); + } + Err(e) => { + warn!( + "Failed to deserialize stored input item: {}. Item: {}", + e, item + ); + } + } + } + } + + // Convert output items from stored output (which is now a JSON array) + if let Some(output_arr) = stored.output.as_array() { + for item in output_arr { + match serde_json::from_value::(item.clone()) { + Ok(output_item) => { + items.push(output_item); + } + Err(e) => { + warn!( + "Failed to deserialize stored output item: {}. Item: {}", + e, item + ); + } + } + } + } + } + conversation_items = Some(items); + modified_request.previous_response_id = None; + } + Err(e) => { + warn!( + "Failed to load previous response chain for {}: {}", + prev_id_str, e + ); + } + } + } + + // Handle conversation by loading conversation history + if let Some(ref conv_id_str) = request.conversation { + let conv_id = ConversationId::from(conv_id_str.as_str()); + + // Check if conversation exists - return error if not found + let conversation = ctx + .conversation_storage + .get_conversation(&conv_id) + .await + .map_err(|e| { + error::internal_error( + "check_conversation_failed", + format!("Failed to check conversation: {}", e), + ) + })?; + + if conversation.is_none() { + return Err(error::not_found( + "conversation_not_found", + format!( + "Conversation '{}' not found. Please create the conversation first using the conversations API.", + conv_id_str + ) + )); + } + + // Load conversation history + const MAX_CONVERSATION_HISTORY_ITEMS: usize = 100; + let params = data_connector::ListParams { + limit: MAX_CONVERSATION_HISTORY_ITEMS, + order: data_connector::SortOrder::Asc, + after: None, + }; + + match ctx + .conversation_item_storage + .list_items(&conv_id, params) + .await + { + Ok(stored_items) => { + let mut items: Vec = Vec::new(); + for item in stored_items.into_iter() { + if item.item_type == "message" { + if let Ok(content_parts) = + serde_json::from_value::>(item.content.clone()) + { + items.push(ResponseInputOutputItem::Message { + id: item.id.0.clone(), + role: item.role.clone().unwrap_or_else(|| "user".to_string()), + content: content_parts, + status: item.status.clone(), + }); + } + } + } + + // Append current request + match &modified_request.input { + ResponseInput::Text(text) => { + items.push(ResponseInputOutputItem::Message { + id: format!("msg_u_{}", conv_id.0), + role: "user".to_string(), + content: vec![ResponseContentPart::InputText { text: text.clone() }], + status: Some("completed".to_string()), + }); + } + ResponseInput::Items(current_items) => { + // Process all item types, converting SimpleInputMessage to Message + for item in current_items.iter() { + let normalized = responses::normalize_input_item(item); + items.push(normalized); + } + } + } + + modified_request.input = ResponseInput::Items(items); + } + Err(e) => { + warn!("Failed to load conversation history: {}", e); + } + } + } + + // If we have conversation_items from previous_response_id, merge them + if let Some(mut items) = conversation_items { + // Append current request + match &modified_request.input { + ResponseInput::Text(text) => { + items.push(ResponseInputOutputItem::Message { + id: format!( + "msg_u_{}", + request + .previous_response_id + .as_ref() + .unwrap_or(&"new".to_string()) + ), + role: "user".to_string(), + content: vec![ResponseContentPart::InputText { text: text.clone() }], + status: Some("completed".to_string()), + }); + } + ResponseInput::Items(current_items) => { + // Process all item types, converting SimpleInputMessage to Message + for item in current_items.iter() { + let normalized = responses::normalize_input_item(item); + items.push(normalized); + } + } + } + + modified_request.input = ResponseInput::Items(items); + } + + debug!( + has_previous_response = request.previous_response_id.is_some(), + has_conversation = request.conversation.is_some(), + "Loaded conversation history" + ); + + Ok(modified_request) +} + +/// Build next request with updated conversation history +pub(super) fn build_next_request( + state: &ToolLoopState, + current_request: &ResponsesRequest, +) -> ResponsesRequest { + // Start with original input + let mut input_items = match &state.original_input { + ResponseInput::Text(text) => vec![ResponseInputOutputItem::Message { + id: format!("msg_u_{}", state.iteration), + role: "user".to_string(), + content: vec![ResponseContentPart::InputText { text: text.clone() }], + status: Some("completed".to_string()), + }], + ResponseInput::Items(items) => items.iter().map(responses::normalize_input_item).collect(), + }; + + // Append all conversation history (function calls and outputs) + input_items.extend_from_slice(&state.conversation_history); + + // Build new request for next iteration + ResponsesRequest { + input: ResponseInput::Items(input_items), + model: current_request.model.clone(), + instructions: current_request.instructions.clone(), + tools: current_request.tools.clone(), + max_output_tokens: current_request.max_output_tokens, + temperature: current_request.temperature, + top_p: current_request.top_p, + stream: current_request.stream, + store: Some(false), // Don't store intermediate responses + background: Some(false), + max_tool_calls: current_request.max_tool_calls, + tool_choice: current_request.tool_choice.clone(), + parallel_tool_calls: current_request.parallel_tool_calls, + previous_response_id: None, + conversation: None, + user: current_request.user.clone(), + metadata: current_request.metadata.clone(), + include: current_request.include.clone(), + reasoning: current_request.reasoning.clone(), + service_tier: current_request.service_tier.clone(), + top_logprobs: current_request.top_logprobs, + truncation: current_request.truncation.clone(), + text: current_request.text.clone(), + request_id: None, + priority: current_request.priority, + frequency_penalty: current_request.frequency_penalty, + presence_penalty: current_request.presence_penalty, + stop: current_request.stop.clone(), + top_k: current_request.top_k, + min_p: current_request.min_p, + repetition_penalty: current_request.repetition_penalty, + } +} diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs index 80c0b98a7..b0b81e044 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/context.rs @@ -1,20 +1,33 @@ -//! Context for /v1/responses endpoint handlers +//! Context and types for /v1/responses endpoint handlers //! //! Bundles all dependencies needed by responses handlers to avoid passing //! 10+ parameters to every function. use std::{collections::HashMap, sync::Arc}; -use tokio::sync::RwLock; +use tokio::{sync::RwLock, task::JoinHandle}; -use super::types::BackgroundTaskInfo; use crate::{ core::WorkerRegistry, data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, + grpc_client::SglangSchedulerClient, mcp::McpManager, routers::grpc::{context::SharedComponents, pipeline::RequestPipeline}, }; +/// Information stored for background tasks to enable end-to-end cancellation +/// +/// This struct enables cancelling both the Rust task AND the Python scheduler processing. +/// The client field is lazily initialized during pipeline execution. +pub struct BackgroundTaskInfo { + /// Tokio task handle for aborting the Rust task + pub handle: JoinHandle<()>, + /// gRPC request_id sent to Python scheduler (chatcmpl-* prefix) + pub grpc_request_id: String, + /// gRPC client for sending abort requests to Python (set after client acquisition) + pub client: Arc>>, +} + /// Context for /v1/responses endpoint /// /// All fields are Arc/shared references, so cloning this context is cheap. diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs index 794f2e327..7ff0aded0 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs @@ -3,16 +3,14 @@ //! # Public API //! //! - `route_responses()` - POST /v1/responses (main entry point) -//! - `get_response_impl()` - GET /v1/responses/{response_id} -//! - `cancel_response_impl()` - POST /v1/responses/{response_id}/cancel //! //! # Architecture //! -//! This module orchestrates all request handling for the /v1/responses endpoint. +//! This module provides the entry point for the /v1/responses endpoint. //! It supports two execution modes: //! -//! 1. **Synchronous** - Returns complete response immediately -//! 2. **Streaming** - Returns SSE stream with real-time events +//! 1. **Synchronous** - Returns complete response immediately (non_streaming.rs) +//! 2. **Streaming** - Returns SSE stream with real-time events (streaming.rs) //! //! Note: Background mode is no longer supported. Requests with background=true //! will be rejected with a 400 error. @@ -21,63 +19,35 @@ //! //! ```text //! route_responses() -//! ├─► route_responses_sync() → route_responses_internal() -//! └─► route_responses_streaming() → convert_chat_stream_to_responses_stream() -//! -//! route_responses_internal() -//! ├─► load_conversation_history() -//! ├─► execute_tool_loop() (if MCP tools) -//! │ └─► pipeline.execute_chat_for_responses() [loop] -//! └─► execute_without_mcp() (if no MCP tools) -//! └─► pipeline.execute_chat_for_responses() +//! ├─► route_responses_sync() → non_streaming::route_responses_internal() +//! └─► route_responses_streaming() +//! ├─► streaming::execute_tool_loop_streaming() (MCP tools) +//! └─► streaming::convert_chat_stream_to_responses_stream() (no MCP) //! ``` use std::sync::Arc; use axum::{ - body::Body, http, response::{IntoResponse, Response}, }; -use bytes::Bytes; -use futures_util::StreamExt; -use serde_json::json; -use tokio::sync::mpsc; -use tracing::{debug, error, warn}; +use tracing::debug; use uuid::Uuid; use super::{ - conversions, - tool_loop::{execute_tool_loop, execute_tool_loop_streaming}, + common::load_conversation_history, context::ResponsesContext, conversions, non_streaming, + streaming, }; use crate::{ - data_connector::{ - self, ConversationId, ConversationItemStorage, ConversationStorage, ResponseId, - ResponseStorage, - }, - protocols::{ - chat::{self, ChatCompletionStreamResponse}, - common::{self}, - responses::{ - self, ResponseContentPart, ResponseInput, ResponseInputOutputItem, ResponseOutputItem, - ResponseReasoningContent, ResponseStatus, ResponsesRequest, ResponsesResponse, - ResponsesUsage, - }, - }, - routers::{ - error, - grpc::common::responses::{ - build_sse_response, ensure_mcp_connection, persist_response_if_needed, - streaming::ResponseStreamEventEmitter, - }, - }, + protocols::responses::ResponsesRequest, + routers::{error, grpc::common::responses::ensure_mcp_connection}, }; /// Main handler for POST /v1/responses /// -/// Validates request, determines execution mode (sync/async/streaming), and delegates +/// Validates request, determines execution mode (sync/streaming), and delegates pub async fn route_responses( - ctx: &super::context::ResponsesContext, + ctx: &ResponsesContext, request: Arc, headers: Option, model_id: Option, @@ -103,86 +73,32 @@ pub async fn route_responses( } // ============================================================================ -// Synchronous Execution +// Synchronous Entry Point // ============================================================================ /// Execute synchronous responses request -/// -/// This is the core execution path that: -/// 1. Loads conversation history / response chain -/// 2. Converts to ChatCompletionRequest -/// 3. Executes chat pipeline -/// 4. Converts back to ResponsesResponse -/// 5. Persists to storage async fn route_responses_sync( - ctx: &super::context::ResponsesContext, + ctx: &ResponsesContext, request: Arc, headers: Option, model_id: Option, response_id: Option, ) -> Response { - match route_responses_internal(ctx, request, headers, model_id, response_id).await { + match non_streaming::route_responses_internal(ctx, request, headers, model_id, response_id) + .await + { Ok(responses_response) => axum::Json(responses_response).into_response(), Err(response) => response, // Already a Response with proper status code } } -/// Internal implementation that returns Result for background task compatibility -async fn route_responses_internal( - ctx: &super::context::ResponsesContext, - request: Arc, - headers: Option, - model_id: Option, - response_id: Option, -) -> Result { - // 1. Load conversation history and build modified request - let modified_request = load_conversation_history(ctx, &request).await?; - - // 2. Check MCP connection and get whether MCP tools are present - let has_mcp_tools = ensure_mcp_connection(&ctx.mcp_manager, request.tools.as_deref()).await?; - - let responses_response = if has_mcp_tools { - debug!("MCP tools detected, using tool loop"); - - // Execute with MCP tool loop - execute_tool_loop( - ctx, - modified_request, - &request, - headers, - model_id, - response_id.clone(), - ) - .await? - } else { - // No MCP tools - execute without MCP (may have function tools or no tools) - execute_without_mcp( - ctx, - &modified_request, - &request, - headers, - model_id, - response_id.clone(), - ) - .await? - }; - - // 5. Persist response to storage if store=true - persist_response_if_needed( - ctx.conversation_storage.clone(), - ctx.conversation_item_storage.clone(), - ctx.response_storage.clone(), - &responses_response, - &request, - ) - .await; - - Ok(responses_response) -} +// ============================================================================ +// Streaming Entry Point +// ============================================================================ /// Execute streaming responses request async fn route_responses_streaming( - ctx: &super::context::ResponsesContext, + ctx: &ResponsesContext, request: Arc, headers: Option, model_id: Option, @@ -203,8 +119,14 @@ async fn route_responses_streaming( if has_mcp_tools { debug!("MCP tools detected in streaming mode, using streaming tool loop"); - return execute_tool_loop_streaming(ctx, modified_request, &request, headers, model_id) - .await; + return streaming::execute_tool_loop_streaming( + ctx, + modified_request, + &request, + headers, + model_id, + ) + .await; } // 3. Convert ResponsesRequest → ChatCompletionRequest @@ -219,594 +141,12 @@ async fn route_responses_streaming( }; // 4. Execute chat pipeline and convert streaming format (no MCP tools) - convert_chat_stream_to_responses_stream(ctx, chat_request, headers, model_id, &request).await -} - -/// Convert chat streaming response to responses streaming format -/// -/// This function: -/// 1. Gets chat SSE stream from pipeline -/// 2. Intercepts and parses each SSE event -/// 3. Converts ChatCompletionStreamResponse → ResponsesResponse delta -/// 4. Accumulates response state for final persistence -/// 5. Emits transformed SSE events in responses format -async fn convert_chat_stream_to_responses_stream( - ctx: &super::context::ResponsesContext, - chat_request: Arc, - headers: Option, - model_id: Option, - original_request: &ResponsesRequest, -) -> Response { - debug!("Converting chat SSE stream to responses SSE format"); - - // Get chat streaming response - let chat_response = ctx - .pipeline - .execute_chat( - chat_request.clone(), - headers, - model_id, - ctx.components.clone(), - ) - .await; - - // Extract body from chat response - let (_parts, body) = chat_response.into_parts(); - - // Create channel for transformed SSE events - let (tx, rx) = mpsc::unbounded_channel::>(); - - // Spawn background task to transform stream - let original_request_clone = original_request.clone(); - let chat_request_clone = chat_request.clone(); - let response_storage = ctx.response_storage.clone(); - let conversation_storage = ctx.conversation_storage.clone(); - let conversation_item_storage = ctx.conversation_item_storage.clone(); - - tokio::spawn(async move { - if let Err(e) = process_and_transform_sse_stream( - body, - original_request_clone, - chat_request_clone, - response_storage, - conversation_storage, - conversation_item_storage, - tx.clone(), - ) - .await - { - warn!("Error transforming SSE stream: {}", e); - let error_event = json!({ - "error": { - "message": e, - "type": "stream_error" - } - }); - let _ = tx.send(Ok(Bytes::from(format!("data: {}\n\n", error_event)))); - } - - // Send final [DONE] event - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); - }); - - // Build SSE response with transformed stream - build_sse_response(rx) -} - -/// Process chat SSE stream and transform to responses format -async fn process_and_transform_sse_stream( - body: Body, - original_request: ResponsesRequest, - _chat_request: Arc, - response_storage: Arc, - conversation_storage: Arc, - conversation_item_storage: Arc, - tx: mpsc::UnboundedSender>, -) -> Result<(), String> { - // Create accumulator for final response - let mut accumulator = StreamingResponseAccumulator::new(&original_request); - - // Create event emitter for OpenAI-compatible streaming - let response_id = format!("resp_{}", Uuid::new_v4()); - let model = original_request.model.clone(); - let created_at = chrono::Utc::now().timestamp() as u64; - let mut event_emitter = ResponseStreamEventEmitter::new(response_id, model, created_at); - event_emitter.set_original_request(original_request.clone()); - - // Emit initial response.created and response.in_progress events - let event = event_emitter.emit_created(); - event_emitter - .send_event(&event, &tx) - .map_err(|_| "Failed to send response.created event".to_string())?; - - let event = event_emitter.emit_in_progress(); - event_emitter - .send_event(&event, &tx) - .map_err(|_| "Failed to send response.in_progress event".to_string())?; - - // Convert body to data stream - let mut stream = body.into_data_stream(); - - // Process stream chunks (each chunk is a complete SSE event) - while let Some(chunk_result) = stream.next().await { - let chunk = chunk_result.map_err(|e| format!("Stream read error: {}", e))?; - - // Convert chunk to string - let event_str = String::from_utf8_lossy(&chunk); - let event = event_str.trim(); - - // Check for end of stream - if event == "data: [DONE]" { - break; - } - - // Parse SSE event (format: "data: {...}\n\n" or "data: {...}") - if let Some(json_str) = event.strip_prefix("data: ") { - let json_str = json_str.trim(); - - // Try to parse as ChatCompletionStreamResponse - match serde_json::from_str::(json_str) { - Ok(chat_chunk) => { - // Update accumulator - accumulator.process_chunk(&chat_chunk); - - // Process chunk through event emitter (emits proper OpenAI events) - event_emitter.process_chunk(&chat_chunk, &tx)?; - } - Err(_) => { - // Not a valid chat chunk - might be error event, pass through - debug!("Non-chunk SSE event, passing through: {}", event); - if tx.send(Ok(Bytes::from(format!("{}\n\n", event)))).is_err() { - return Err("Client disconnected".to_string()); - } - } - } - } - } - - // Emit final response.completed event with accumulated usage - let usage_json = accumulator.usage.as_ref().map(|u| { - let mut usage_obj = json!({ - "input_tokens": u.prompt_tokens, - "output_tokens": u.completion_tokens, - "total_tokens": u.total_tokens - }); - - // Include reasoning_tokens if present - if let Some(details) = &u.completion_tokens_details { - if let Some(reasoning_tokens) = details.reasoning_tokens { - usage_obj["output_tokens_details"] = - json!({ "reasoning_tokens": reasoning_tokens }); - } - } - - usage_obj - }); - - let completed_event = event_emitter.emit_completed(usage_json.as_ref()); - event_emitter.send_event(&completed_event, &tx)?; - - // Finalize and persist accumulated response - let final_response = accumulator.finalize(); - persist_response_if_needed( - conversation_storage, - conversation_item_storage, - response_storage, - &final_response, - &original_request, + streaming::convert_chat_stream_to_responses_stream( + ctx, + chat_request, + headers, + model_id, + &request, ) - .await; - - Ok(()) -} - -/// Response accumulator for streaming responses -struct StreamingResponseAccumulator { - // Response metadata - response_id: String, - model: String, - created_at: i64, - - // Accumulated content - content_buffer: String, - reasoning_buffer: String, - tool_calls: Vec, - - // Completion state - finish_reason: Option, - usage: Option, - - // Original request for final response construction - original_request: ResponsesRequest, -} - -impl StreamingResponseAccumulator { - fn new(original_request: &ResponsesRequest) -> Self { - Self { - response_id: String::new(), - model: String::new(), - created_at: 0, - content_buffer: String::new(), - reasoning_buffer: String::new(), - tool_calls: Vec::new(), - finish_reason: None, - usage: None, - original_request: original_request.clone(), - } - } - - fn process_chunk(&mut self, chunk: &ChatCompletionStreamResponse) { - // Initialize metadata on first chunk - if self.response_id.is_empty() { - self.response_id = chunk.id.clone(); - self.model = chunk.model.clone(); - self.created_at = chunk.created as i64; - } - - // Process first choice (responses API doesn't support n>1) - if let Some(choice) = chunk.choices.first() { - // Accumulate content - if let Some(content) = &choice.delta.content { - self.content_buffer.push_str(content); - } - - // Accumulate reasoning - if let Some(reasoning) = &choice.delta.reasoning_content { - self.reasoning_buffer.push_str(reasoning); - } - - // Process tool call deltas - if let Some(tool_call_deltas) = &choice.delta.tool_calls { - for delta in tool_call_deltas { - // Use index directly (it's a u32, not Option) - let index = delta.index as usize; - - // Ensure we have enough tool calls - while self.tool_calls.len() <= index { - self.tool_calls.push(ResponseOutputItem::FunctionToolCall { - id: String::new(), - call_id: String::new(), - name: String::new(), - arguments: String::new(), - output: None, - status: "in_progress".to_string(), - }); - } - - // Update the tool call at this index - if let ResponseOutputItem::FunctionToolCall { - id, - name, - arguments, - .. - } = &mut self.tool_calls[index] - { - if let Some(delta_id) = &delta.id { - id.push_str(delta_id); - } - if let Some(function) = &delta.function { - if let Some(delta_name) = &function.name { - name.push_str(delta_name); - } - if let Some(delta_args) = &function.arguments { - arguments.push_str(delta_args); - } - } - } - } - } - - // Update finish reason - if let Some(reason) = &choice.finish_reason { - self.finish_reason = Some(reason.clone()); - } - } - - // Update usage - if let Some(usage) = &chunk.usage { - self.usage = Some(usage.clone()); - } - } - - fn finalize(self) -> ResponsesResponse { - let mut output: Vec = Vec::new(); - - // Add message content if present - if !self.content_buffer.is_empty() { - output.push(ResponseOutputItem::Message { - id: format!("msg_{}", self.response_id), - role: "assistant".to_string(), - content: vec![ResponseContentPart::OutputText { - text: self.content_buffer, - annotations: vec![], - logprobs: None, - }], - status: "completed".to_string(), - }); - } - - // Add reasoning if present - if !self.reasoning_buffer.is_empty() { - output.push(ResponseOutputItem::Reasoning { - id: format!("reasoning_{}", self.response_id), - summary: vec![], - content: vec![ResponseReasoningContent::ReasoningText { - text: self.reasoning_buffer, - }], - status: Some("completed".to_string()), - }); - } - - // Add tool calls - output.extend(self.tool_calls); - - // Determine final status - let status = match self.finish_reason.as_deref() { - Some("stop") | Some("length") => ResponseStatus::Completed, - Some("tool_calls") => ResponseStatus::InProgress, - Some("failed") | Some("error") => ResponseStatus::Failed, - _ => ResponseStatus::Completed, - }; - - // Convert usage - let usage = self.usage.as_ref().map(|u| { - let usage_info = common::UsageInfo { - prompt_tokens: u.prompt_tokens, - completion_tokens: u.completion_tokens, - total_tokens: u.total_tokens, - reasoning_tokens: u - .completion_tokens_details - .as_ref() - .and_then(|d| d.reasoning_tokens), - prompt_tokens_details: None, - }; - ResponsesUsage::Classic(usage_info) - }); - - ResponsesResponse::builder(&self.response_id, &self.model) - .copy_from_request(&self.original_request) - .created_at(self.created_at) - .status(status) - .output(output) - .maybe_usage(usage) - .build() - } -} - -// ============================================================================ -// Helper Functions -// ============================================================================ - -/// Execute request without MCP tool loop (simple pipeline execution) -async fn execute_without_mcp( - ctx: &super::context::ResponsesContext, - modified_request: &ResponsesRequest, - original_request: &ResponsesRequest, - headers: Option, - model_id: Option, - response_id: Option, -) -> Result { - // Convert ResponsesRequest → ChatCompletionRequest - let chat_request = conversions::responses_to_chat(modified_request).map_err(|e| { - error!( - function = "execute_without_mcp", - error = %e, - "Failed to convert ResponsesRequest to ChatCompletionRequest" - ); - error::bad_request( - "convert_request_failed", - format!("Failed to convert request: {}", e), - ) - })?; - - // Execute chat pipeline (errors already have proper HTTP status codes) - let chat_response = ctx - .pipeline - .execute_chat_for_responses( - Arc::new(chat_request), - headers, - model_id, - ctx.components.clone(), - ) - .await?; // Preserve the Response error as-is - - // Convert ChatCompletionResponse → ResponsesResponse - conversions::chat_to_responses(&chat_response, original_request, response_id).map_err(|e| { - error!( - function = "execute_without_mcp", - error = %e, - "Failed to convert ChatCompletionResponse to ResponsesResponse" - ); - error::internal_error( - "convert_to_responses_format_failed", - format!("Failed to convert to responses format: {}", e), - ) - }) -} - -/// Load conversation history and response chains, returning modified request -async fn load_conversation_history( - ctx: &super::context::ResponsesContext, - request: &ResponsesRequest, -) -> Result { - let mut modified_request = request.clone(); - let mut conversation_items: Option> = None; - - // Handle previous_response_id by loading response chain - if let Some(ref prev_id_str) = modified_request.previous_response_id { - let prev_id = ResponseId::from(prev_id_str.as_str()); - match ctx - .response_storage - .get_response_chain(&prev_id, None) - .await - { - Ok(chain) => { - let mut items = Vec::new(); - for stored in chain.responses.iter() { - // Convert input items from stored input (which is now a JSON array) - if let Some(input_arr) = stored.input.as_array() { - for item in input_arr { - match serde_json::from_value::(item.clone()) { - Ok(input_item) => { - items.push(input_item); - } - Err(e) => { - warn!( - "Failed to deserialize stored input item: {}. Item: {}", - e, item - ); - } - } - } - } - - // Convert output items from stored output (which is now a JSON array) - if let Some(output_arr) = stored.output.as_array() { - for item in output_arr { - match serde_json::from_value::(item.clone()) { - Ok(output_item) => { - items.push(output_item); - } - Err(e) => { - warn!( - "Failed to deserialize stored output item: {}. Item: {}", - e, item - ); - } - } - } - } - } - conversation_items = Some(items); - modified_request.previous_response_id = None; - } - Err(e) => { - warn!( - "Failed to load previous response chain for {}: {}", - prev_id_str, e - ); - } - } - } - - // Handle conversation by loading conversation history - if let Some(ref conv_id_str) = request.conversation { - let conv_id = ConversationId::from(conv_id_str.as_str()); - - // Check if conversation exists - return error if not found - let conversation = ctx - .conversation_storage - .get_conversation(&conv_id) - .await - .map_err(|e| { - error!( - function = "load_conversation_history", - conversation_id = %conv_id_str, - error = %e, - "Failed to check conversation existence in storage" - ); - error::internal_error( - "check_conversation_failed", - format!("Failed to check conversation: {}", e), - ) - })?; - - if conversation.is_none() { - return Err(error::not_found( - "conversation_not_found", - format!( - "Conversation '{}' not found. Please create the conversation first using the conversations API.", - conv_id_str - ) - )); - } - - // Load conversation history - const MAX_CONVERSATION_HISTORY_ITEMS: usize = 100; - let params = data_connector::ListParams { - limit: MAX_CONVERSATION_HISTORY_ITEMS, - order: data_connector::SortOrder::Asc, - after: None, - }; - - match ctx - .conversation_item_storage - .list_items(&conv_id, params) - .await - { - Ok(stored_items) => { - let mut items: Vec = Vec::new(); - for item in stored_items.into_iter() { - if item.item_type == "message" { - if let Ok(content_parts) = - serde_json::from_value::>(item.content.clone()) - { - items.push(ResponseInputOutputItem::Message { - id: item.id.0.clone(), - role: item.role.clone().unwrap_or_else(|| "user".to_string()), - content: content_parts, - status: item.status.clone(), - }); - } - } - } - - // Append current request - match &modified_request.input { - ResponseInput::Text(text) => { - items.push(ResponseInputOutputItem::Message { - id: format!("msg_u_{}", conv_id.0), - role: "user".to_string(), - content: vec![ResponseContentPart::InputText { text: text.clone() }], - status: Some("completed".to_string()), - }); - } - ResponseInput::Items(current_items) => { - // Process all item types, converting SimpleInputMessage to Message - for item in current_items.iter() { - let normalized = responses::normalize_input_item(item); - items.push(normalized); - } - } - } - - modified_request.input = ResponseInput::Items(items); - } - Err(e) => { - warn!("Failed to load conversation history: {}", e); - } - } - } - - // If we have conversation_items from previous_response_id, merge them - if let Some(mut items) = conversation_items { - // Append current request - match &modified_request.input { - ResponseInput::Text(text) => { - items.push(ResponseInputOutputItem::Message { - id: format!( - "msg_u_{}", - request - .previous_response_id - .as_ref() - .unwrap_or(&"new".to_string()) - ), - role: "user".to_string(), - content: vec![ResponseContentPart::InputText { text: text.clone() }], - status: Some("completed".to_string()), - }); - } - ResponseInput::Items(current_items) => { - // Process all item types, converting SimpleInputMessage to Message - for item in current_items.iter() { - let normalized = responses::normalize_input_item(item); - items.push(normalized); - } - } - } - - modified_request.input = ResponseInput::Items(items); - } - - Ok(modified_request) + .await } diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/mod.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/mod.rs index ce67c0191..e4b55c20f 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/mod.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/mod.rs @@ -1,20 +1,23 @@ //! Regular gRPC Router `/v1/responses` endpoint implementation //! -//! This module handles all responses-specific logic for the regular (non-Harmony) pipeline including: -//! - Request validation -//! - Conversation history and response chain loading -//! - Streaming support -//! - MCP tool loop wrapper -//! - Response persistence +//! This module handles all responses-specific logic for the regular (non-Harmony) pipeline. +//! +//! ## Architecture +//! +//! - `handlers` - Entry points: route_responses (thin dispatcher) +//! - `non_streaming` - Non-streaming execution with MCP tool loop +//! - `streaming` - Streaming execution with MCP tool loop +//! - `common` - Shared helpers: ToolLoopState, tool preparation, MCP metadata builders +//! - `conversions` - Request/response conversion between Responses and Chat formats +//! - `context` - ResponsesContext and BackgroundTaskInfo -// Module declarations -pub mod context; +mod common; +mod context; mod conversions; mod handlers; -pub mod tool_loop; -pub mod types; +mod non_streaming; +mod streaming; // Public exports -pub use context::ResponsesContext; +pub use context::{BackgroundTaskInfo, ResponsesContext}; pub use handlers::route_responses; -pub use types::BackgroundTaskInfo; diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs new file mode 100644 index 000000000..d8a83c204 --- /dev/null +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs @@ -0,0 +1,417 @@ +//! Non-streaming execution for Regular Responses API +//! +//! This module handles non-streaming request execution: +//! - `route_responses_internal` - Core execution orchestration +//! - `execute_tool_loop` - MCP tool loop execution +//! - `execute_without_mcp` - Simple pipeline execution without MCP + +use std::{sync::Arc, time::Instant}; + +use axum::response::Response; +use serde_json::json; +use tracing::{debug, error, trace, warn}; + +use super::{ + common::{ + build_mcp_list_tools_item, build_next_request, convert_mcp_tools_to_chat_tools, + extract_all_tool_calls_from_chat, load_conversation_history, prepare_chat_tools_and_choice, + ToolLoopState, + }, + context::ResponsesContext, + conversions, +}; +use crate::{ + observability::metrics::{metrics_labels, Metrics}, + protocols::responses::{ResponseStatus, ResponsesRequest, ResponsesResponse}, + routers::{ + error, + grpc::common::responses::{ensure_mcp_connection, persist_response_if_needed}, + mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS}, + }, +}; + +/// Internal implementation for non-streaming responses +/// +/// This is the core execution path that: +/// 1. Loads conversation history / response chain +/// 2. Checks for MCP tools +/// 3. Executes with or without MCP tool loop +/// 4. Persists to storage +pub(super) async fn route_responses_internal( + ctx: &ResponsesContext, + request: Arc, + headers: Option, + model_id: Option, + response_id: Option, +) -> Result { + // 1. Load conversation history and build modified request + let modified_request = load_conversation_history(ctx, &request).await?; + + // 2. Check MCP connection and get whether MCP tools are present + let has_mcp_tools = ensure_mcp_connection(&ctx.mcp_manager, request.tools.as_deref()).await?; + + let responses_response = if has_mcp_tools { + debug!("MCP tools detected, using tool loop"); + + // Execute with MCP tool loop + execute_tool_loop( + ctx, + modified_request, + &request, + headers, + model_id, + response_id.clone(), + ) + .await? + } else { + // No MCP tools - execute without MCP (may have function tools or no tools) + execute_without_mcp( + ctx, + &modified_request, + &request, + headers, + model_id, + response_id.clone(), + ) + .await? + }; + + // 5. Persist response to storage if store=true + persist_response_if_needed( + ctx.conversation_storage.clone(), + ctx.conversation_item_storage.clone(), + ctx.response_storage.clone(), + &responses_response, + &request, + ) + .await; + + Ok(responses_response) +} + +/// Execute request without MCP tool loop (simple pipeline execution) +pub(super) async fn execute_without_mcp( + ctx: &ResponsesContext, + modified_request: &ResponsesRequest, + original_request: &ResponsesRequest, + headers: Option, + model_id: Option, + response_id: Option, +) -> Result { + // Convert ResponsesRequest → ChatCompletionRequest + let chat_request = conversions::responses_to_chat(modified_request).map_err(|e| { + error!( + function = "execute_without_mcp", + error = %e, + "Failed to convert ResponsesRequest to ChatCompletionRequest" + ); + error::bad_request( + "convert_request_failed", + format!("Failed to convert request: {}", e), + ) + })?; + + // Execute chat pipeline (errors already have proper HTTP status codes) + let chat_response = ctx + .pipeline + .execute_chat_for_responses( + Arc::new(chat_request), + headers, + model_id, + ctx.components.clone(), + ) + .await?; // Preserve the Response error as-is + + // Convert ChatCompletionResponse → ResponsesResponse + conversions::chat_to_responses(&chat_response, original_request, response_id).map_err(|e| { + error!( + function = "execute_without_mcp", + error = %e, + "Failed to convert ChatCompletionResponse to ResponsesResponse" + ); + error::internal_error( + "convert_to_responses_format_failed", + format!("Failed to convert to responses format: {}", e), + ) + }) +} + +/// Execute the MCP tool calling loop +/// +/// This wraps pipeline.execute_chat_for_responses() in a loop that: +/// 1. Executes the chat pipeline +/// 2. Checks if response has tool calls +/// 3. If yes, executes MCP tools and builds resume request +/// 4. Repeats until no more tool calls or limit reached +pub(super) async fn execute_tool_loop( + ctx: &ResponsesContext, + mut current_request: ResponsesRequest, + original_request: &ResponsesRequest, + headers: Option, + model_id: Option, + response_id: Option, +) -> Result { + // Get server label from original request tools + let server_label = extract_server_label(original_request.tools.as_deref(), "request-mcp"); + + let mut state = ToolLoopState::new(original_request.input.clone(), server_label.clone()); + + // Configuration: max iterations as safety limit + let max_tool_calls = original_request.max_tool_calls.map(|n| n as usize); + + trace!( + "Starting MCP tool loop: server_label={}, max_tool_calls={:?}, max_iterations={}", + server_label, + max_tool_calls, + DEFAULT_MAX_ITERATIONS + ); + + // Get MCP tools and convert to chat format (do this once before loop) + let mcp_tools = ctx.mcp_manager.list_tools(); + let mcp_chat_tools = convert_mcp_tools_to_chat_tools(&mcp_tools); + trace!( + "Converted {} MCP tools to chat format", + mcp_chat_tools.len() + ); + + loop { + // Convert to chat request + let mut chat_request = conversions::responses_to_chat(¤t_request).map_err(|e| { + error!( + function = "tool_loop", + iteration = state.iteration, + error = %e, + "Failed to convert ResponsesRequest to ChatCompletionRequest in tool loop" + ); + error::bad_request( + "convert_request_failed", + format!("Failed to convert request: {}", e), + ) + })?; + + // Prepare tools and tool_choice for this iteration + prepare_chat_tools_and_choice(&mut chat_request, &mcp_chat_tools, state.iteration); + + // Execute chat pipeline (errors already have proper HTTP status codes) + let chat_response = ctx + .pipeline + .execute_chat_for_responses( + Arc::new(chat_request), + headers.clone(), + model_id.clone(), + ctx.components.clone(), + ) + .await?; + + // Check for function calls (extract all for parallel execution) + let tool_calls = extract_all_tool_calls_from_chat(&chat_response); + + if !tool_calls.is_empty() { + state.iteration += 1; + + // Record tool loop iteration metric + Metrics::record_mcp_tool_iteration(¤t_request.model); + + trace!( + "Tool loop iteration {}: found {} tool call(s)", + state.iteration, + tool_calls.len() + ); + + // Separate MCP and function tool calls + let mcp_tool_names: std::collections::HashSet<&str> = + mcp_tools.iter().map(|t| t.name.as_ref()).collect(); + let (mcp_tool_calls, function_tool_calls): (Vec<_>, Vec<_>) = tool_calls + .into_iter() + .partition(|(_, tool_name, _)| mcp_tool_names.contains(tool_name.as_str())); + + trace!( + "Separated tool calls: {} MCP, {} function", + mcp_tool_calls.len(), + function_tool_calls.len() + ); + + // If ANY tool call is a function tool, return to caller immediately + if !function_tool_calls.is_empty() { + // Convert chat response to responses format (includes all tool calls) + let responses_response = conversions::chat_to_responses( + &chat_response, + original_request, + response_id.clone(), + ) + .map_err(|e| { + error!( + function = "tool_loop", + iteration = state.iteration, + error = %e, + context = "function_tool_calls", + "Failed to convert ChatCompletionResponse to ResponsesResponse" + ); + error::internal_error( + "convert_to_responses_format_failed", + format!("Failed to convert to responses format: {}", e), + ) + })?; + + // Return response with function tool calls to caller + return Ok(responses_response); + } + + // All MCP tools - check combined limit BEFORE executing + let effective_limit = match max_tool_calls { + Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS), + None => DEFAULT_MAX_ITERATIONS, + }; + + if state.total_calls + mcp_tool_calls.len() > effective_limit { + warn!( + "Reached tool call limit: {} + {} > {} (max_tool_calls={:?}, safety_limit={})", + state.total_calls, + mcp_tool_calls.len(), + effective_limit, + max_tool_calls, + DEFAULT_MAX_ITERATIONS + ); + + // Convert chat response to responses format and mark as incomplete + let mut responses_response = conversions::chat_to_responses( + &chat_response, + original_request, + response_id.clone(), + ) + .map_err(|e| { + error!( + function = "tool_loop", + iteration = state.iteration, + error = %e, + context = "max_tool_calls_limit", + "Failed to convert ChatCompletionResponse to ResponsesResponse" + ); + error::internal_error( + "convert_to_responses_format_failed", + format!("Failed to convert to responses format: {}", e), + ) + })?; + + // Mark as completed but with incomplete details + responses_response.status = ResponseStatus::Completed; + responses_response.incomplete_details = Some(json!({ "reason": "max_tool_calls" })); + + return Ok(responses_response); + } + + // Execute all MCP tools + for (call_id, tool_name, args_json_str) in mcp_tool_calls { + trace!( + "Calling MCP tool '{}' (call_id: {}) with args: {}", + tool_name, + call_id, + args_json_str + ); + + let tool_start = Instant::now(); + let (output_str, success, error) = match ctx + .mcp_manager + .call_tool(tool_name.as_str(), args_json_str.as_str()) + .await + { + Ok(result) => match serde_json::to_string(&result) { + Ok(output) => (output, true, None), + Err(e) => { + let err = format!("Failed to serialize tool result: {}", e); + warn!("{}", err); + let error_json = json!({ "error": &err }).to_string(); + (error_json, false, Some(err)) + } + }, + Err(err) => { + let err_str = format!("tool call failed: {}", err); + warn!("Tool execution failed: {}", err_str); + // Return error as output, let model decide how to proceed + let error_json = json!({ "error": &err_str }).to_string(); + (error_json, false, Some(err_str)) + } + }; + let tool_duration = tool_start.elapsed(); + + // Record MCP tool metrics + Metrics::record_mcp_tool_duration( + ¤t_request.model, + &tool_name, + tool_duration, + ); + Metrics::record_mcp_tool_call( + ¤t_request.model, + &tool_name, + if success { + metrics_labels::RESULT_SUCCESS + } else { + metrics_labels::RESULT_ERROR + }, + ); + + // Record the call in state + state.record_call( + call_id, + tool_name, + args_json_str, + output_str, + success, + error, + ); + + // Increment total calls counter + state.total_calls += 1; + } + + // Build resume request with conversation history + current_request = build_next_request(&state, ¤t_request); + + // Continue to next iteration + } else { + // No more tool calls, we're done + trace!( + "Tool loop completed: {} iterations, {} total calls", + state.iteration, + state.total_calls + ); + + // Convert final chat response to responses format + let mut responses_response = conversions::chat_to_responses( + &chat_response, + original_request, + response_id.clone(), + ) + .map_err(|e| { + error!( + function = "tool_loop", + iteration = state.iteration, + error = %e, + context = "final_response", + "Failed to convert ChatCompletionResponse to ResponsesResponse" + ); + error::internal_error( + "convert_to_responses_format_failed", + format!("Failed to convert to responses format: {}", e), + ) + })?; + + // Inject MCP metadata into output + if state.total_calls > 0 { + // Prepend mcp_list_tools item + let mcp_list_tools = build_mcp_list_tools_item(&ctx.mcp_manager, &server_label); + responses_response.output.insert(0, mcp_list_tools); + + // Append all mcp_call items at the end + responses_response.output.extend(state.mcp_call_items); + + trace!( + "Injected MCP metadata: 1 mcp_list_tools + {} mcp_call items", + state.total_calls + ); + } + + return Ok(responses_response); + } + } +} diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs similarity index 54% rename from sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs rename to sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs index 6bda1af0f..211080cf4 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs @@ -1,4 +1,9 @@ -//! MCP tool loop execution for /v1/responses endpoint +//! Streaming execution for Regular Responses API +//! +//! This module handles streaming request execution: +//! - `execute_tool_loop_streaming` - MCP tool loop with streaming +//! - `convert_chat_stream_to_responses_stream` - Non-MCP streaming conversion +//! - Streaming accumulators for response building use std::{ collections::HashMap, @@ -8,7 +13,7 @@ use std::{ use axum::{ body::Body, - http::{header, StatusCode}, + http::{self, header, StatusCode}, response::Response, }; use bytes::Bytes; @@ -16,530 +21,403 @@ use futures_util::StreamExt; use serde_json::{json, Value}; use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; -use tracing::{error, trace, warn}; +use tracing::{debug, trace, warn}; use uuid::Uuid; -use super::conversions; +use super::{ + common::{ + build_next_request, convert_mcp_tools_to_chat_tools, extract_all_tool_calls_from_chat, + prepare_chat_tools_and_choice, ToolLoopState, + }, + context::ResponsesContext, + conversions, +}; use crate::{ - mcp::{self, McpManager}, + data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, observability::metrics::{metrics_labels, Metrics}, protocols::{ chat::{ ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionStreamResponse, }, - common::{Function, FunctionCallResponse, Tool, ToolCall, ToolChoice, ToolChoiceValue}, + common::{FunctionCallResponse, ToolCall, Usage, UsageInfo}, responses::{ - self, McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem, - ResponseOutputItem, ResponseStatus, ResponsesRequest, ResponsesResponse, + ResponseContentPart, ResponseOutputItem, ResponseReasoningContent, ResponseStatus, + ResponsesRequest, ResponsesResponse, ResponsesUsage, }, }, routers::{ - error, - grpc::common::responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, + grpc::common::responses::{ + build_sse_response, persist_response_if_needed, + streaming::{OutputItemType, ResponseStreamEventEmitter}, + }, mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS}, }, }; -/// Merge function tools from request with MCP tools and set tool_choice based on iteration -fn prepare_chat_tools_and_choice( - chat_request: &mut ChatCompletionRequest, - mcp_chat_tools: &[Tool], - iteration: usize, -) { - // Merge function tools from request with MCP tools - let mut all_tools = chat_request.tools.clone().unwrap_or_default(); - all_tools.extend(mcp_chat_tools.iter().cloned()); - chat_request.tools = Some(all_tools); - - // Set tool_choice based on iteration - // - Iteration 0: Use user's tool_choice or default to auto - // - Iteration 1+: Always use auto to avoid infinite loops - chat_request.tool_choice = if iteration == 0 { - chat_request - .tool_choice - .clone() - .or(Some(ToolChoice::Value(ToolChoiceValue::Auto))) - } else { - Some(ToolChoice::Value(ToolChoiceValue::Auto)) - }; -} - -/// Extract all tool calls from chat response (for parallel tool call support) -fn extract_all_tool_calls_from_chat( - response: &ChatCompletionResponse, -) -> Vec<(String, String, String)> { - // Check if response has choices with tool calls - let Some(choice) = response.choices.first() else { - return Vec::new(); - }; - let message = &choice.message; - - // Look for tool_calls in the message - if let Some(tool_calls) = &message.tool_calls { - tool_calls - .iter() - .map(|tool_call| { - ( - tool_call.id.clone(), - tool_call.function.name.clone(), - tool_call - .function - .arguments - .clone() - .unwrap_or_else(|| "{}".to_string()), - ) - }) - .collect() - } else { - Vec::new() - } -} - -/// State for tracking multi-turn tool calling loop -struct ToolLoopState { - iteration: usize, - total_calls: usize, - conversation_history: Vec, - original_input: ResponseInput, - mcp_call_items: Vec, - server_label: String, -} - -impl ToolLoopState { - fn new(original_input: ResponseInput, server_label: String) -> Self { - Self { - iteration: 0, - total_calls: 0, - conversation_history: Vec::new(), - original_input, - mcp_call_items: Vec::new(), - server_label, - } - } - - fn record_call( - &mut self, - call_id: String, - tool_name: String, - args_json_str: String, - output_str: String, - success: bool, - error: Option, - ) { - // Add function_tool_call item with both arguments and output - self.conversation_history - .push(ResponseInputOutputItem::FunctionToolCall { - id: call_id.clone(), - call_id: call_id.clone(), - name: tool_name.clone(), - arguments: args_json_str.clone(), - output: Some(output_str.clone()), - status: Some("completed".to_string()), - }); - - // Add mcp_call output item for metadata - let mcp_call = build_mcp_call_item( - &tool_name, - &args_json_str, - &output_str, - &self.server_label, - success, - error.as_deref(), - ); - self.mcp_call_items.push(mcp_call); - } -} - // ============================================================================ -// MCP Metadata Builders +// Non-MCP Streaming Path // ============================================================================ -/// Generate unique ID for MCP items -fn generate_mcp_id(prefix: &str) -> String { - format!("{}_{}", prefix, Uuid::new_v4()) -} - -/// Build mcp_list_tools output item -fn build_mcp_list_tools_item(mcp: &Arc, server_label: &str) -> ResponseOutputItem { - let tools = mcp.list_tools(); - let tools_info: Vec = tools - .iter() - .map(|t| McpToolInfo { - name: t.name.to_string(), - description: t.description.as_ref().map(|d| d.to_string()), - input_schema: Value::Object((*t.input_schema).clone()), - annotations: Some(json!({ - "read_only": false - })), - }) - .collect(); - - ResponseOutputItem::McpListTools { - id: generate_mcp_id("mcpl"), - server_label: server_label.to_string(), - tools: tools_info, - } -} - -/// Build mcp_call output item -fn build_mcp_call_item( - tool_name: &str, - arguments: &str, - output: &str, - server_label: &str, - success: bool, - error: Option<&str>, -) -> ResponseOutputItem { - ResponseOutputItem::McpCall { - id: generate_mcp_id("mcp"), - status: if success { "completed" } else { "failed" }.to_string(), - approval_request_id: None, - arguments: arguments.to_string(), - error: error.map(|e| e.to_string()), - name: tool_name.to_string(), - output: output.to_string(), - server_label: server_label.to_string(), - } -} - -/// Execute the MCP tool calling loop +/// Convert chat streaming response to responses streaming format /// -/// This wraps pipeline.execute_chat_for_responses() in a loop that: -/// 1. Executes the chat pipeline -/// 2. Checks if response has tool calls -/// 3. If yes, executes MCP tools and builds resume request -/// 4. Repeats until no more tool calls or limit reached -pub(super) async fn execute_tool_loop( - ctx: &super::context::ResponsesContext, - mut current_request: ResponsesRequest, - original_request: &ResponsesRequest, +/// This function: +/// 1. Gets chat SSE stream from pipeline +/// 2. Intercepts and parses each SSE event +/// 3. Converts ChatCompletionStreamResponse → ResponsesResponse delta +/// 4. Accumulates response state for final persistence +/// 5. Emits transformed SSE events in responses format +pub(super) async fn convert_chat_stream_to_responses_stream( + ctx: &ResponsesContext, + chat_request: Arc, headers: Option, model_id: Option, - response_id: Option, -) -> Result { - // Get server label from original request tools - let server_label = extract_server_label(original_request.tools.as_deref(), "request-mcp"); + original_request: &ResponsesRequest, +) -> Response { + debug!("Converting chat SSE stream to responses SSE format"); - let mut state = ToolLoopState::new(original_request.input.clone(), server_label.clone()); + // Get chat streaming response + let chat_response = ctx + .pipeline + .execute_chat( + chat_request.clone(), + headers, + model_id, + ctx.components.clone(), + ) + .await; - // Configuration: max iterations as safety limit - let max_tool_calls = original_request.max_tool_calls.map(|n| n as usize); + // Extract body from chat response + let (_parts, body) = chat_response.into_parts(); - trace!( - "Starting MCP tool loop: server_label={}, max_tool_calls={:?}, max_iterations={}", - server_label, - max_tool_calls, - DEFAULT_MAX_ITERATIONS - ); + // Create channel for transformed SSE events + let (tx, rx) = mpsc::unbounded_channel::>(); - // Get MCP tools and convert to chat format (do this once before loop) - let mcp_tools = ctx.mcp_manager.list_tools(); - let mcp_chat_tools = convert_mcp_tools_to_chat_tools(&mcp_tools); - trace!( - "Converted {} MCP tools to chat format", - mcp_chat_tools.len() - ); + // Spawn background task to transform stream + let original_request_clone = original_request.clone(); + let response_storage = ctx.response_storage.clone(); + let conversation_storage = ctx.conversation_storage.clone(); + let conversation_item_storage = ctx.conversation_item_storage.clone(); - loop { - // Convert to chat request - let mut chat_request = conversions::responses_to_chat(¤t_request).map_err(|e| { - error!( - function = "tool_loop", - iteration = state.iteration, - error = %e, - "Failed to convert ResponsesRequest to ChatCompletionRequest in tool loop" - ); - error::bad_request( - "convert_request_failed", - format!("Failed to convert request: {}", e), - ) - })?; - - // Prepare tools and tool_choice for this iteration - prepare_chat_tools_and_choice(&mut chat_request, &mcp_chat_tools, state.iteration); - - // Execute chat pipeline (errors already have proper HTTP status codes) - let chat_response = ctx - .pipeline - .execute_chat_for_responses( - Arc::new(chat_request), - headers.clone(), - model_id.clone(), - ctx.components.clone(), - ) - .await?; - - // Check for function calls (extract all for parallel execution) - let tool_calls = extract_all_tool_calls_from_chat(&chat_response); - - if !tool_calls.is_empty() { - state.iteration += 1; - - // Record tool loop iteration metric - Metrics::record_mcp_tool_iteration(¤t_request.model); - - trace!( - "Tool loop iteration {}: found {} tool call(s)", - state.iteration, - tool_calls.len() - ); - - // Separate MCP and function tool calls - let mcp_tool_names: std::collections::HashSet<&str> = - mcp_tools.iter().map(|t| t.name.as_ref()).collect(); - let (mcp_tool_calls, function_tool_calls): (Vec<_>, Vec<_>) = tool_calls - .into_iter() - .partition(|(_, tool_name, _)| mcp_tool_names.contains(tool_name.as_str())); - - trace!( - "Separated tool calls: {} MCP, {} function", - mcp_tool_calls.len(), - function_tool_calls.len() - ); - - // If ANY tool call is a function tool, return to caller immediately - if !function_tool_calls.is_empty() { - // Convert chat response to responses format (includes all tool calls) - let responses_response = conversions::chat_to_responses( - &chat_response, - original_request, - response_id.clone(), - ) - .map_err(|e| { - error!( - function = "tool_loop", - iteration = state.iteration, - error = %e, - context = "function_tool_calls", - "Failed to convert ChatCompletionResponse to ResponsesResponse" - ); - error::internal_error( - "convert_to_responses_format_failed", - format!("Failed to convert to responses format: {}", e), - ) - })?; - - // Return response with function tool calls to caller - return Ok(responses_response); - } - - // All MCP tools - check combined limit BEFORE executing - let effective_limit = match max_tool_calls { - Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS), - None => DEFAULT_MAX_ITERATIONS, - }; - - if state.total_calls + mcp_tool_calls.len() > effective_limit { - warn!( - "Reached tool call limit: {} + {} > {} (max_tool_calls={:?}, safety_limit={})", - state.total_calls, - mcp_tool_calls.len(), - effective_limit, - max_tool_calls, - DEFAULT_MAX_ITERATIONS - ); - - // Convert chat response to responses format and mark as incomplete - let mut responses_response = conversions::chat_to_responses( - &chat_response, - original_request, - response_id.clone(), - ) - .map_err(|e| { - error!( - function = "tool_loop", - iteration = state.iteration, - error = %e, - context = "max_tool_calls_limit", - "Failed to convert ChatCompletionResponse to ResponsesResponse" - ); - error::internal_error( - "convert_to_responses_format_failed", - format!("Failed to convert to responses format: {}", e), - ) - })?; - - // Mark as completed but with incomplete details - responses_response.status = ResponseStatus::Completed; - responses_response.incomplete_details = Some(json!({ "reason": "max_tool_calls" })); - - return Ok(responses_response); - } - - // Execute all MCP tools - for (call_id, tool_name, args_json_str) in mcp_tool_calls { - trace!( - "Calling MCP tool '{}' (call_id: {}) with args: {}", - tool_name, - call_id, - args_json_str - ); - - let tool_start = Instant::now(); - let (output_str, success, error) = match ctx - .mcp_manager - .call_tool(tool_name.as_str(), args_json_str.as_str()) - .await - { - Ok(result) => match serde_json::to_string(&result) { - Ok(output) => (output, true, None), - Err(e) => { - let err = format!("Failed to serialize tool result: {}", e); - warn!("{}", err); - let error_json = json!({ "error": &err }).to_string(); - (error_json, false, Some(err)) - } - }, - Err(err) => { - let err_str = format!("tool call failed: {}", err); - warn!("Tool execution failed: {}", err_str); - // Return error as output, let model decide how to proceed - let error_json = json!({ "error": &err_str }).to_string(); - (error_json, false, Some(err_str)) - } - }; - let tool_duration = tool_start.elapsed(); - - // Record MCP tool metrics - Metrics::record_mcp_tool_duration( - ¤t_request.model, - &tool_name, - tool_duration, - ); - Metrics::record_mcp_tool_call( - ¤t_request.model, - &tool_name, - if success { - metrics_labels::RESULT_SUCCESS - } else { - metrics_labels::RESULT_ERROR - }, - ); - - // Record the call in state - state.record_call( - call_id, - tool_name, - args_json_str, - output_str, - success, - error, - ); - - // Increment total calls counter - state.total_calls += 1; - } - - // Build resume request with conversation history - // Start with original input - let mut input_items = match &state.original_input { - ResponseInput::Text(text) => vec![ResponseInputOutputItem::Message { - id: format!("msg_u_{}", state.iteration), - role: "user".to_string(), - content: vec![ResponseContentPart::InputText { text: text.clone() }], - status: Some("completed".to_string()), - }], - ResponseInput::Items(items) => { - items.iter().map(responses::normalize_input_item).collect() + tokio::spawn(async move { + if let Err(e) = process_and_transform_sse_stream( + body, + original_request_clone, + response_storage, + conversation_storage, + conversation_item_storage, + tx.clone(), + ) + .await + { + warn!("Error transforming SSE stream: {}", e); + let error_event = json!({ + "error": { + "message": e, + "type": "stream_error" } - }; + }); + let _ = tx.send(Ok(Bytes::from(format!("data: {}\n\n", error_event)))); + } - // Append all conversation history (function calls and outputs) - input_items.extend_from_slice(&state.conversation_history); + // Send final [DONE] event + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + }); - // Build new request for next iteration - current_request = ResponsesRequest { - input: ResponseInput::Items(input_items), - model: current_request.model.clone(), - instructions: current_request.instructions.clone(), - tools: current_request.tools.clone(), - max_output_tokens: current_request.max_output_tokens, - temperature: current_request.temperature, - top_p: current_request.top_p, - stream: Some(false), // Always non-streaming in tool loop - store: Some(false), // Don't store intermediate responses - background: Some(false), - max_tool_calls: current_request.max_tool_calls, - tool_choice: current_request.tool_choice.clone(), - parallel_tool_calls: current_request.parallel_tool_calls, - previous_response_id: None, - conversation: None, - user: current_request.user.clone(), - metadata: current_request.metadata.clone(), - // Additional fields from ResponsesRequest - include: current_request.include.clone(), - reasoning: current_request.reasoning.clone(), - service_tier: current_request.service_tier.clone(), - top_logprobs: current_request.top_logprobs, - truncation: current_request.truncation.clone(), - text: current_request.text.clone(), - request_id: None, - priority: current_request.priority, - frequency_penalty: current_request.frequency_penalty, - presence_penalty: current_request.presence_penalty, - stop: current_request.stop.clone(), - top_k: current_request.top_k, - min_p: current_request.min_p, - repetition_penalty: current_request.repetition_penalty, - }; + // Build SSE response with transformed stream + build_sse_response(rx) +} - // Continue to next iteration - } else { - // No more tool calls, we're done - trace!( - "Tool loop completed: {} iterations, {} total calls", - state.iteration, - state.total_calls - ); +/// Process chat SSE stream and transform to responses format +async fn process_and_transform_sse_stream( + body: Body, + original_request: ResponsesRequest, + response_storage: Arc, + conversation_storage: Arc, + conversation_item_storage: Arc, + tx: mpsc::UnboundedSender>, +) -> Result<(), String> { + // Create accumulator for final response + let mut accumulator = StreamingResponseAccumulator::new(&original_request); - // Convert final chat response to responses format - let mut responses_response = conversions::chat_to_responses( - &chat_response, - original_request, - response_id.clone(), - ) - .map_err(|e| { - error!( - function = "tool_loop", - iteration = state.iteration, - error = %e, - context = "final_response", - "Failed to convert ChatCompletionResponse to ResponsesResponse" - ); - error::internal_error( - "convert_to_responses_format_failed", - format!("Failed to convert to responses format: {}", e), - ) - })?; + // Create event emitter for OpenAI-compatible streaming + let response_id = format!("resp_{}", Uuid::new_v4()); + let model = original_request.model.clone(); + let created_at = chrono::Utc::now().timestamp() as u64; + let mut event_emitter = ResponseStreamEventEmitter::new(response_id, model, created_at); + event_emitter.set_original_request(original_request.clone()); - // Inject MCP metadata into output - if state.total_calls > 0 { - // Prepend mcp_list_tools item - let mcp_list_tools = build_mcp_list_tools_item(&ctx.mcp_manager, &server_label); - responses_response.output.insert(0, mcp_list_tools); + // Emit initial response.created and response.in_progress events + let event = event_emitter.emit_created(); + event_emitter + .send_event(&event, &tx) + .map_err(|_| "Failed to send response.created event".to_string())?; - // Append all mcp_call items at the end - responses_response.output.extend(state.mcp_call_items); + let event = event_emitter.emit_in_progress(); + event_emitter + .send_event(&event, &tx) + .map_err(|_| "Failed to send response.in_progress event".to_string())?; - trace!( - "Injected MCP metadata: 1 mcp_list_tools + {} mcp_call items", - state.total_calls - ); + // Convert body to data stream + let mut stream = body.into_data_stream(); + + // Process stream chunks (each chunk is a complete SSE event) + while let Some(chunk_result) = stream.next().await { + let chunk = chunk_result.map_err(|e| format!("Stream read error: {}", e))?; + + // Convert chunk to string + let event_str = String::from_utf8_lossy(&chunk); + let event = event_str.trim(); + + // Check for end of stream + if event == "data: [DONE]" { + break; + } + + // Parse SSE event (format: "data: {...}\n\n" or "data: {...}") + if let Some(json_str) = event.strip_prefix("data: ") { + let json_str = json_str.trim(); + + // Try to parse as ChatCompletionStreamResponse + match serde_json::from_str::(json_str) { + Ok(chat_chunk) => { + // Update accumulator + accumulator.process_chunk(&chat_chunk); + + // Process chunk through event emitter (emits proper OpenAI events) + event_emitter.process_chunk(&chat_chunk, &tx)?; + } + Err(_) => { + // Not a valid chat chunk - might be error event, pass through + debug!("Non-chunk SSE event, passing through: {}", event); + if tx.send(Ok(Bytes::from(format!("{}\n\n", event)))).is_err() { + return Err("Client disconnected".to_string()); + } + } } - - return Ok(responses_response); } } + + // Emit final response.completed event with accumulated usage + let usage_json = accumulator.usage.as_ref().map(|u| { + let mut usage_obj = json!({ + "input_tokens": u.prompt_tokens, + "output_tokens": u.completion_tokens, + "total_tokens": u.total_tokens + }); + + // Include reasoning_tokens if present + if let Some(details) = &u.completion_tokens_details { + if let Some(reasoning_tokens) = details.reasoning_tokens { + usage_obj["output_tokens_details"] = + json!({ "reasoning_tokens": reasoning_tokens }); + } + } + + usage_obj + }); + + let completed_event = event_emitter.emit_completed(usage_json.as_ref()); + event_emitter.send_event(&completed_event, &tx)?; + + // Finalize and persist accumulated response + let final_response = accumulator.finalize(); + persist_response_if_needed( + conversation_storage, + conversation_item_storage, + response_storage, + &final_response, + &original_request, + ) + .await; + + Ok(()) } +/// Response accumulator for streaming responses (non-MCP path) +struct StreamingResponseAccumulator { + // Response metadata + response_id: String, + model: String, + created_at: i64, + + // Accumulated content + content_buffer: String, + reasoning_buffer: String, + tool_calls: Vec, + + // Completion state + finish_reason: Option, + usage: Option, + + // Original request for final response construction + original_request: ResponsesRequest, +} + +impl StreamingResponseAccumulator { + fn new(original_request: &ResponsesRequest) -> Self { + Self { + response_id: String::new(), + model: String::new(), + created_at: 0, + content_buffer: String::new(), + reasoning_buffer: String::new(), + tool_calls: Vec::new(), + finish_reason: None, + usage: None, + original_request: original_request.clone(), + } + } + + fn process_chunk(&mut self, chunk: &ChatCompletionStreamResponse) { + // Initialize metadata on first chunk + if self.response_id.is_empty() { + self.response_id = chunk.id.clone(); + self.model = chunk.model.clone(); + self.created_at = chunk.created as i64; + } + + // Process first choice (responses API doesn't support n>1) + if let Some(choice) = chunk.choices.first() { + // Accumulate content + if let Some(content) = &choice.delta.content { + self.content_buffer.push_str(content); + } + + // Accumulate reasoning + if let Some(reasoning) = &choice.delta.reasoning_content { + self.reasoning_buffer.push_str(reasoning); + } + + // Process tool call deltas + if let Some(tool_call_deltas) = &choice.delta.tool_calls { + for delta in tool_call_deltas { + // Use index directly (it's a u32, not Option) + let index = delta.index as usize; + + // Ensure we have enough tool calls + while self.tool_calls.len() <= index { + self.tool_calls.push(ResponseOutputItem::FunctionToolCall { + id: String::new(), + call_id: String::new(), + name: String::new(), + arguments: String::new(), + output: None, + status: "in_progress".to_string(), + }); + } + + // Update the tool call at this index + if let ResponseOutputItem::FunctionToolCall { + id, + name, + arguments, + .. + } = &mut self.tool_calls[index] + { + if let Some(delta_id) = &delta.id { + id.push_str(delta_id); + } + if let Some(function) = &delta.function { + if let Some(delta_name) = &function.name { + name.push_str(delta_name); + } + if let Some(delta_args) = &function.arguments { + arguments.push_str(delta_args); + } + } + } + } + } + + // Update finish reason + if let Some(reason) = &choice.finish_reason { + self.finish_reason = Some(reason.clone()); + } + } + + // Update usage + if let Some(usage) = &chunk.usage { + self.usage = Some(usage.clone()); + } + } + + fn finalize(self) -> ResponsesResponse { + let mut output: Vec = Vec::new(); + + // Add message content if present + if !self.content_buffer.is_empty() { + output.push(ResponseOutputItem::Message { + id: format!("msg_{}", self.response_id), + role: "assistant".to_string(), + content: vec![ResponseContentPart::OutputText { + text: self.content_buffer, + annotations: vec![], + logprobs: None, + }], + status: "completed".to_string(), + }); + } + + // Add reasoning if present + if !self.reasoning_buffer.is_empty() { + output.push(ResponseOutputItem::Reasoning { + id: format!("reasoning_{}", self.response_id), + summary: vec![], + content: vec![ResponseReasoningContent::ReasoningText { + text: self.reasoning_buffer, + }], + status: Some("completed".to_string()), + }); + } + + // Add tool calls + output.extend(self.tool_calls); + + // Determine final status + let status = match self.finish_reason.as_deref() { + Some("stop") | Some("length") => ResponseStatus::Completed, + Some("tool_calls") => ResponseStatus::InProgress, + Some("failed") | Some("error") => ResponseStatus::Failed, + _ => ResponseStatus::Completed, + }; + + // Convert usage + let usage = self.usage.as_ref().map(|u| { + let usage_info = UsageInfo { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + reasoning_tokens: u + .completion_tokens_details + .as_ref() + .and_then(|d| d.reasoning_tokens), + prompt_tokens_details: None, + }; + ResponsesUsage::Classic(usage_info) + }); + + ResponsesResponse::builder(&self.response_id, &self.model) + .copy_from_request(&self.original_request) + .created_at(self.created_at) + .status(status) + .output(output) + .maybe_usage(usage) + .build() + } +} + +// ============================================================================ +// MCP Streaming Path +// ============================================================================ + /// Execute MCP tool loop with streaming support /// /// This streams each iteration's response to the client while accumulating /// to check for tool calls. If tool calls are found, executes them and /// continues with the next streaming iteration. pub(super) async fn execute_tool_loop_streaming( - ctx: &super::context::ResponsesContext, + ctx: &ResponsesContext, current_request: ResponsesRequest, original_request: &ResponsesRequest, headers: Option, @@ -606,7 +484,7 @@ pub(super) async fn execute_tool_loop_streaming( /// Internal streaming tool loop implementation async fn execute_tool_loop_streaming_internal( - ctx: &super::context::ResponsesContext, + ctx: &ResponsesContext, mut current_request: ResponsesRequest, original_request: &ResponsesRequest, headers: Option, @@ -1004,53 +882,7 @@ async fn execute_tool_loop_streaming_internal( } // Build next request with conversation history - let mut input_items = match &state.original_input { - ResponseInput::Text(text) => vec![ResponseInputOutputItem::Message { - id: format!("msg_u_{}", state.iteration), - role: "user".to_string(), - content: vec![ResponseContentPart::InputText { text: text.clone() }], - status: Some("completed".to_string()), - }], - ResponseInput::Items(items) => { - items.iter().map(responses::normalize_input_item).collect() - } - }; - - input_items.extend_from_slice(&state.conversation_history); - - current_request = ResponsesRequest { - input: ResponseInput::Items(input_items), - model: current_request.model.clone(), - instructions: current_request.instructions.clone(), - tools: current_request.tools.clone(), - max_output_tokens: current_request.max_output_tokens, - temperature: current_request.temperature, - top_p: current_request.top_p, - stream: Some(true), - store: Some(false), - background: Some(false), - max_tool_calls: current_request.max_tool_calls, - tool_choice: current_request.tool_choice.clone(), - parallel_tool_calls: current_request.parallel_tool_calls, - previous_response_id: None, - conversation: None, - user: current_request.user.clone(), - metadata: current_request.metadata.clone(), - include: current_request.include.clone(), - reasoning: current_request.reasoning.clone(), - service_tier: current_request.service_tier.clone(), - top_logprobs: current_request.top_logprobs, - truncation: current_request.truncation.clone(), - text: current_request.text.clone(), - request_id: None, - priority: current_request.priority, - frequency_penalty: current_request.frequency_penalty, - presence_penalty: current_request.presence_penalty, - stop: current_request.stop.clone(), - top_k: current_request.top_k, - min_p: current_request.min_p, - repetition_penalty: current_request.repetition_penalty, - }; + current_request = build_next_request(&state, ¤t_request); continue; } @@ -1091,22 +923,6 @@ async fn execute_tool_loop_streaming_internal( Ok(()) } -/// Convert MCP tools to Chat API tool format -fn convert_mcp_tools_to_chat_tools(mcp_tools: &[mcp::Tool]) -> Vec { - mcp_tools - .iter() - .map(|tool_info| Tool { - tool_type: "function".to_string(), - function: Function { - name: tool_info.name.to_string(), - description: tool_info.description.as_ref().map(|d| d.to_string()), - parameters: Value::Object((*tool_info.input_schema).clone()), - strict: None, - }, - }) - .collect() -} - /// Convert chat stream to Responses API events while accumulating for tool call detection async fn convert_and_accumulate_stream( body: Body, @@ -1147,8 +963,10 @@ struct ChatResponseAccumulator { id: String, model: String, content: String, + reasoning_content: Option, tool_calls: HashMap, finish_reason: Option, + usage: Option, } impl ChatResponseAccumulator { @@ -1157,8 +975,10 @@ impl ChatResponseAccumulator { id: String::new(), model: String::new(), content: String::new(), + reasoning_content: None, tool_calls: HashMap::new(), finish_reason: None, + usage: None, } } @@ -1176,6 +996,13 @@ impl ChatResponseAccumulator { self.content.push_str(content); } + // Accumulate reasoning content + if let Some(reasoning) = &choice.delta.reasoning_content { + self.reasoning_content + .get_or_insert_with(String::new) + .push_str(reasoning); + } + // Accumulate tool calls if let Some(tool_call_deltas) = &choice.delta.tool_calls { for delta in tool_call_deltas { @@ -1210,6 +1037,11 @@ impl ChatResponseAccumulator { self.finish_reason = Some(reason.clone()); } } + + // Update usage + if let Some(usage) = &chunk.usage { + self.usage = Some(usage.clone()); + } } fn finalize(self) -> ChatCompletionResponse { @@ -1232,13 +1064,14 @@ impl ChatResponseAccumulator { } else { Some(tool_calls) }, - reasoning_content: None, + reasoning_content: self.reasoning_content, }, finish_reason: self.finish_reason, logprobs: None, matched_stop: None, hidden_states: None, }]) + .maybe_usage(self.usage) .build() } } diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/types.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/types.rs deleted file mode 100644 index fb6921b2d..000000000 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/types.rs +++ /dev/null @@ -1,18 +0,0 @@ -//! Type definitions for /v1/responses endpoint - -use std::sync::Arc; - -use tokio::{sync::RwLock, task::JoinHandle}; - -/// Information stored for background tasks to enable end-to-end cancellation -/// -/// This struct enables cancelling both the Rust task AND the Python scheduler processing. -/// The client field is lazily initialized during pipeline execution. -pub struct BackgroundTaskInfo { - /// Tokio task handle for aborting the Rust task - pub handle: JoinHandle<()>, - /// gRPC request_id sent to Python scheduler (chatcmpl-* prefix) - pub grpc_request_id: String, - /// gRPC client for sending abort requests to Python (set after client acquisition) - pub client: Arc>>, -}