diff --git a/sgl-model-gateway/src/routers/conversations/handlers.rs b/sgl-model-gateway/src/routers/conversations/handlers.rs index 91755e857..07babe4f8 100644 --- a/sgl-model-gateway/src/routers/conversations/handlers.rs +++ b/sgl-model-gateway/src/routers/conversations/handlers.rs @@ -1,4 +1,4 @@ -//! Conversation CRUD handlers - shared across routers +//! Conversation CRUD handlers for the /v1/conversations API - shared across routers use std::sync::Arc; @@ -9,11 +9,15 @@ use axum::{ }; use chrono::Utc; use serde_json::{json, Value}; -use tracing::{debug, info, warn}; +use tracing::{info, warn}; -use crate::data_connector::{ - Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, - ConversationStorage, ListParams, NewConversation, NewConversationItem, SortOrder, +use crate::{ + data_connector::{ + Conversation, ConversationId, ConversationItem, ConversationItemId, + ConversationItemStorage, ConversationStorage, ListParams, NewConversation, + NewConversationItem, SortOrder, + }, + routers::persistence_utils::item_to_json, }; // ============================================================================ @@ -551,47 +555,6 @@ pub async fn delete_conversation_item( } } -// ============================================================================ -// Item Creation Helper -// ============================================================================ - -pub async fn create_and_link_item( - item_storage: &Arc, - conv_id_opt: Option<&ConversationId>, - mut new_item: NewConversationItem, -) -> Result<(), String> { - if new_item.status.is_none() { - new_item.status = Some("completed".to_string()); - } - - let created = item_storage - .create_item(new_item) - .await - .map_err(|e| format!("Failed to create item: {e}"))?; - - if let Some(conv_id) = conv_id_opt { - item_storage - .link_item(conv_id, &created.id, Utc::now()) - .await - .map_err(|e| format!("Failed to link item: {e}"))?; - - debug!( - conversation_id = %conv_id.0, - item_id = %created.id.0, - item_type = %created.item_type, - "Persisted conversation item and link" - ); - } else { - debug!( - item_id = %created.id.0, - item_type = %created.item_type, - "Persisted conversation item (no conversation link)" - ); - } - - Ok(()) -} - // ============================================================================ // Parsing and Serialization // ============================================================================ @@ -656,60 +619,6 @@ fn parse_item_from_value( )) } -/// Field mappings for item types that store data in content -const ITEM_TYPE_FIELDS: &[(&str, &[&str])] = &[ - ( - "mcp_call", - &[ - "name", - "arguments", - "output", - "server_label", - "approval_request_id", - "error", - ], - ), - ("mcp_list_tools", &["tools", "server_label"]), - ("function_call", &["call_id", "name", "arguments", "output"]), - ("function_call_output", &["call_id", "output"]), -]; - -pub fn item_to_json(item: &ConversationItem) -> Value { - let mut obj = serde_json::Map::new(); - obj.insert("id".to_string(), json!(item.id.0)); - obj.insert("type".to_string(), json!(item.item_type)); - - if let Some(role) = &item.role { - obj.insert("role".to_string(), json!(role)); - } - - // Find field mappings for this item type - let fields = ITEM_TYPE_FIELDS - .iter() - .find(|(t, _)| *t == item.item_type) - .map(|(_, fields)| *fields); - - if let Some(fields) = fields { - // Extract specific fields from content - if let Some(content_obj) = item.content.as_object() { - for field in fields { - if let Some(value) = content_obj.get(*field) { - obj.insert((*field).to_string(), value.clone()); - } - } - } - } else { - // Default: include content as-is - obj.insert("content".to_string(), item.content.clone()); - } - - if let Some(status) = &item.status { - obj.insert("status".to_string(), json!(status)); - } - - Value::Object(obj) -} - pub fn conversation_to_json(conversation: &Conversation) -> Value { let mut obj = json!({ "id": conversation.id.0, diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs b/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs index 0d74ee743..b44066400 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs @@ -15,8 +15,7 @@ use crate::{ responses::{ResponseTool, ResponseToolType, ResponsesRequest, ResponsesResponse}, }, routers::{ - error, - openai::{conversations::persist_conversation_items, mcp::ensure_request_mcp_client}, + error, mcp_utils::ensure_request_mcp_client, persistence_utils::persist_conversation_items, }, }; diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses.rs index 1f6549729..f41d3eb00 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses.rs @@ -66,12 +66,10 @@ use crate::{ harmony::{processor::ResponsesIterationResult, streaming::HarmonyStreamingProcessor}, pipeline::RequestPipeline, }, + mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS}, }, }; -/// Maximum number of tool execution iterations to prevent infinite loops -const MAX_TOOL_ITERATIONS: usize = 10; - /// Record of a single MCP tool call execution /// /// Stores metadata needed to build mcp_call output items for Responses API format @@ -300,7 +298,7 @@ async fn execute_with_mcp_loop( let mut iteration_count = 0; // Extract server_label from request tools - let server_label = extract_mcp_server_label(current_request.tools.as_deref()); + let server_label = extract_server_label(current_request.tools.as_deref(), "sglang-mcp"); let mut mcp_tracking = McpCallTracking::new(server_label.clone()); // Extract user's max_tool_calls limit (if set) @@ -329,16 +327,19 @@ async fn execute_with_mcp_loop( Metrics::record_mcp_tool_iteration(¤t_request.model); // Safety check: prevent infinite loops - if iteration_count > MAX_TOOL_ITERATIONS { + if iteration_count > DEFAULT_MAX_ITERATIONS { error!( function = "execute_with_mcp_loop", iteration_count = iteration_count, - max_iterations = MAX_TOOL_ITERATIONS, + max_iterations = DEFAULT_MAX_ITERATIONS, "Maximum tool iterations exceeded" ); return Err(error::internal_error( "tool_iterations_exceeded", - format!("Maximum tool iterations ({}) exceeded", MAX_TOOL_ITERATIONS), + format!( + "Maximum tool iterations ({}) exceeded", + DEFAULT_MAX_ITERATIONS + ), )); } @@ -390,8 +391,8 @@ async fn execute_with_mcp_loop( // Check combined limit (user's max_tool_calls vs safety limit) let effective_limit = match max_tool_calls { - Some(user_max) => user_max.min(MAX_TOOL_ITERATIONS), - None => MAX_TOOL_ITERATIONS, + Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS), + None => DEFAULT_MAX_ITERATIONS, }; // Check if we would exceed the limit with these new MCP tool calls @@ -658,7 +659,7 @@ async fn execute_mcp_tool_loop_streaming( tx: &mpsc::UnboundedSender>, ) { // Extract server_label from request tools - let server_label = extract_mcp_server_label(current_request.tools.as_deref()); + let server_label = extract_server_label(current_request.tools.as_deref(), "sglang-mcp"); // Set server label in emitter for MCP call items emitter.set_mcp_server_label(server_label.clone()); @@ -773,9 +774,12 @@ async fn execute_mcp_tool_loop_streaming( Metrics::record_mcp_tool_iteration(¤t_request.model); // Safety check: prevent infinite loops - if iteration_count > MAX_TOOL_ITERATIONS { + if iteration_count > DEFAULT_MAX_ITERATIONS { emitter.emit_error( - &format!("Maximum tool iterations ({}) exceeded", MAX_TOOL_ITERATIONS), + &format!( + "Maximum tool iterations ({}) exceeded", + DEFAULT_MAX_ITERATIONS + ), Some("max_iterations_exceeded"), tx, ); @@ -852,8 +856,8 @@ async fn execute_mcp_tool_loop_streaming( // Check combined limit (user's max_tool_calls vs safety limit) let effective_limit = match max_tool_calls { - Some(user_max) => user_max.min(MAX_TOOL_ITERATIONS), - None => MAX_TOOL_ITERATIONS, + Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS), + None => DEFAULT_MAX_ITERATIONS, }; // Check if we would exceed the limit with these new MCP tool calls @@ -1546,24 +1550,6 @@ fn inject_mcp_metadata( response.output.extend(mcp_call_items); } -/// Extract MCP server label from request tools -/// -/// Searches for the first MCP tool in the tools array and returns its server_label. -/// Falls back to "sglang-mcp" if no MCP tool with server_label is found. -fn extract_mcp_server_label(tools: Option<&[ResponseTool]>) -> String { - tools - .and_then(|tools| { - tools.iter().find_map(|tool| { - if matches!(tool.r#type, ResponseToolType::Mcp) { - tool.server_label.clone() - } else { - None - } - }) - }) - .unwrap_or_else(|| "sglang-mcp".to_string()) -} - /// Load previous conversation messages from storage /// /// If the request has `previous_response_id`, loads the response chain from storage diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs index 0f2df521f..6bda1af0f 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs @@ -31,13 +31,13 @@ use crate::{ common::{Function, FunctionCallResponse, Tool, ToolCall, ToolChoice, ToolChoiceValue}, responses::{ self, McpToolInfo, ResponseContentPart, ResponseInput, ResponseInputOutputItem, - ResponseOutputItem, ResponseStatus, ResponseToolType, ResponsesRequest, - ResponsesResponse, + ResponseOutputItem, ResponseStatus, ResponsesRequest, ResponsesResponse, }, }, routers::{ error, grpc::common::responses::streaming::{OutputItemType, ResponseStreamEventEmitter}, + mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS}, }, }; @@ -219,28 +219,18 @@ pub(super) async fn execute_tool_loop( response_id: Option, ) -> Result { // Get server label from original request tools - let server_label = original_request - .tools - .as_ref() - .and_then(|tools| { - tools - .iter() - .find(|t| matches!(t.r#type, ResponseToolType::Mcp)) - .and_then(|t| t.server_label.clone()) - }) - .unwrap_or_else(|| "request-mcp".to_string()); + 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 - const MAX_ITERATIONS: usize = 10; 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, - MAX_ITERATIONS + DEFAULT_MAX_ITERATIONS ); // Get MCP tools and convert to chat format (do this once before loop) @@ -336,8 +326,8 @@ pub(super) async fn execute_tool_loop( // All MCP tools - check combined limit BEFORE executing let effective_limit = match max_tool_calls { - Some(user_max) => user_max.min(MAX_ITERATIONS), - None => MAX_ITERATIONS, + Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS), + None => DEFAULT_MAX_ITERATIONS, }; if state.total_calls + mcp_tool_calls.len() > effective_limit { @@ -347,7 +337,7 @@ pub(super) async fn execute_tool_loop( mcp_tool_calls.len(), effective_limit, max_tool_calls, - MAX_ITERATIONS + DEFAULT_MAX_ITERATIONS ); // Convert chat response to responses format and mark as incomplete @@ -624,18 +614,8 @@ async fn execute_tool_loop_streaming_internal( tx: mpsc::UnboundedSender>, ) -> Result<(), String> { // Extract server label from original request tools - let server_label = original_request - .tools - .as_ref() - .and_then(|tools| { - tools - .iter() - .find(|t| matches!(t.r#type, ResponseToolType::Mcp)) - .and_then(|t| t.server_label.clone()) - }) - .unwrap_or_else(|| "request-mcp".to_string()); + let server_label = extract_server_label(original_request.tools.as_deref(), "request-mcp"); - const MAX_ITERATIONS: usize = 10; let mut state = ToolLoopState::new(original_request.input.clone(), server_label.clone()); let max_tool_calls = original_request.max_tool_calls.map(|n| n as usize); @@ -672,10 +652,10 @@ async fn execute_tool_loop_streaming_internal( // Record tool loop iteration metric Metrics::record_mcp_tool_iteration(&model); - if state.iteration > MAX_ITERATIONS { + if state.iteration > DEFAULT_MAX_ITERATIONS { return Err(format!( "Tool loop exceeded maximum iterations ({})", - MAX_ITERATIONS + DEFAULT_MAX_ITERATIONS )); } @@ -784,8 +764,8 @@ async fn execute_tool_loop_streaming_internal( // Check combined limit (only count MCP tools since function tools will be returned) let effective_limit = match max_tool_calls { - Some(user_max) => user_max.min(MAX_ITERATIONS), - None => MAX_ITERATIONS, + Some(user_max) => user_max.min(DEFAULT_MAX_ITERATIONS), + None => DEFAULT_MAX_ITERATIONS, }; if state.total_calls + mcp_tool_calls.len() > effective_limit { @@ -795,7 +775,7 @@ async fn execute_tool_loop_streaming_internal( mcp_tool_calls.len(), effective_limit, max_tool_calls, - MAX_ITERATIONS + DEFAULT_MAX_ITERATIONS ); break; } diff --git a/sgl-model-gateway/src/routers/grpc/router.rs b/sgl-model-gateway/src/routers/grpc/router.rs index 3eded9041..3cb508196 100644 --- a/sgl-model-gateway/src/routers/grpc/router.rs +++ b/sgl-model-gateway/src/routers/grpc/router.rs @@ -37,7 +37,6 @@ use crate::{ /// gRPC router implementation for SGLang #[derive(Clone)] -#[allow(dead_code)] pub struct GrpcRouter { worker_registry: Arc, pipeline: RequestPipeline, diff --git a/sgl-model-gateway/src/routers/mcp_utils.rs b/sgl-model-gateway/src/routers/mcp_utils.rs new file mode 100644 index 000000000..e388febd5 --- /dev/null +++ b/sgl-model-gateway/src/routers/mcp_utils.rs @@ -0,0 +1,136 @@ +//! Shared MCP utilities for routers. +//! +//! This module provides shared MCP-related functionality that can be +//! used across different router implementations (OpenAI, gRPC regular, gRPC harmony). + +use std::sync::Arc; + +use tracing::warn; + +use crate::{ + mcp::{McpManager, McpServerConfig, McpTransport}, + protocols::responses::{ResponseTool, ResponseToolType}, +}; + +// ============================================================================ +// Constants +// ============================================================================ + +/// Default maximum tool loop iterations (safety limit). +/// +/// Used as fallback when user doesn't specify `max_tool_calls`. +/// All routers use this same value. +pub const DEFAULT_MAX_ITERATIONS: usize = 10; + +// ============================================================================ +// Configuration +// ============================================================================ + +/// Configuration for MCP tool calling loops. +/// +/// Provides a common structure for loop configuration across routers. +#[derive(Debug, Clone)] +pub struct McpLoopConfig { + /// Maximum iterations as safety limit (default: DEFAULT_MAX_ITERATIONS). + /// Prevents infinite loops when max_tool_calls is not set by user. + pub max_iterations: usize, +} + +impl Default for McpLoopConfig { + fn default() -> Self { + Self { + max_iterations: DEFAULT_MAX_ITERATIONS, + } + } +} + +// ============================================================================ +// Helper Functions +// ============================================================================ + +/// Extract MCP server label from request tools. +/// +/// Searches for the first MCP tool in the tools array and returns its server_label. +/// Falls back to a default value if no MCP tool with server_label is found. +pub fn extract_server_label(tools: Option<&[ResponseTool]>, default_label: &str) -> String { + tools + .and_then(|tools| { + tools.iter().find_map(|tool| { + if matches!(tool.r#type, ResponseToolType::Mcp) { + tool.server_label.clone() + } else { + None + } + }) + }) + .unwrap_or_else(|| default_label.to_string()) +} + +// ============================================================================ +// MCP Connection +// ============================================================================ + +/// Ensure MCP client is connected for request-level MCP tools. +/// +/// This function extracts MCP server configuration from request tools (server_url, authorization) +/// and ensures a client connection is established via the connection pool. +/// +/// Returns `Some(())` if a dynamic MCP tool was found and client was created/retrieved, +/// `None` if no MCP tools with server_url were found or connection failed. +pub async fn ensure_request_mcp_client( + mcp_manager: &Arc, + tools: &[ResponseTool], +) -> Option<()> { + // Find an MCP tool with a server_url + let tool = tools + .iter() + .find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some())?; + + let server_url = tool.server_url.as_ref()?.trim().to_string(); + + // Validate URL scheme + if !(server_url.starts_with("http://") || server_url.starts_with("https://")) { + warn!( + "Ignoring MCP server_url with unsupported scheme: {}", + server_url + ); + return None; + } + + // Extract server label and auth token + let name = tool + .server_label + .clone() + .unwrap_or_else(|| "request-mcp".to_string()); + let token = tool.authorization.clone(); + + // Determine transport type based on URL pattern + let transport = if server_url.contains("/sse") { + McpTransport::Sse { + url: server_url.clone(), + token, + } + } else { + McpTransport::Streamable { + url: server_url.clone(), + token, + } + }; + + // Create server config + let server_config = McpServerConfig { + name, + transport, + proxy: None, + required: false, + }; + + // Use get_or_create_client to establish connection + match mcp_manager.get_or_create_client(server_config).await { + Ok(_client) => Some(()), + Err(err) => { + warn!("Failed to get/create MCP connection: {}", err); + None + } + } +} diff --git a/sgl-model-gateway/src/routers/mod.rs b/sgl-model-gateway/src/routers/mod.rs index 725107258..7301ce2d6 100644 --- a/sgl-model-gateway/src/routers/mod.rs +++ b/sgl-model-gateway/src/routers/mod.rs @@ -26,8 +26,10 @@ pub mod factory; pub mod grpc; pub mod header_utils; pub mod http; +pub mod mcp_utils; pub mod openai; pub mod parse; +pub mod persistence_utils; pub mod router_manager; pub mod tokenize; diff --git a/sgl-model-gateway/src/routers/openai/mcp.rs b/sgl-model-gateway/src/routers/openai/mcp.rs index 6c2dd6b07..47a2b1dab 100644 --- a/sgl-model-gateway/src/routers/openai/mcp.rs +++ b/sgl-model-gateway/src/routers/openai/mcp.rs @@ -20,30 +20,18 @@ use crate::{ mcp, protocols::{ event_types::{is_function_call_type, ItemType, McpEvent, OutputItemEvent}, - responses::{generate_id, ResponseInput, ResponseTool, ResponseToolType, ResponsesRequest}, + responses::{generate_id, ResponseInput, ResponsesRequest}, + }, + routers::{ + header_utils::apply_request_headers, + mcp_utils::{extract_server_label, McpLoopConfig}, }, - routers::header_utils::apply_request_headers, }; // ============================================================================ // Configuration and State Types // ============================================================================ -/// Configuration for MCP tool calling loops -#[allow(dead_code)] -#[derive(Debug, Clone)] -pub(crate) struct McpLoopConfig { - /// Maximum iterations as safety limit (internal only, default: 10) - /// Prevents infinite loops when max_tool_calls is not set - pub max_iterations: usize, -} - -impl Default for McpLoopConfig { - fn default() -> Self { - Self { max_iterations: 10 } - } -} - /// State for tracking multi-turn tool calling loop pub(crate) struct ToolLoopState { /// Current iteration number (starts at 0, increments with each tool call) @@ -126,69 +114,6 @@ impl FunctionCallInProgress { } } -// ============================================================================ -// MCP Manager Integration -// ============================================================================ - -/// Ensure a dynamic MCP client exists for request-scoped tools. -/// -/// This function parses request tools to extract MCP server configuration, -/// then ensures a dynamic client exists in the McpManager via `get_or_create_client()`. -/// The McpManager itself is returned (cloned Arc) for convenience, though the main -/// purpose is the side effect of registering the dynamic client. -/// -/// Returns Some(manager) if a dynamic MCP tool was found and client was created/retrieved, -/// None if no MCP tools were found or connection failed. -pub async fn ensure_request_mcp_client( - mcp_manager: &Arc, - tools: &[ResponseTool], -) -> Option> { - let tool = tools - .iter() - .find(|t| matches!(t.r#type, ResponseToolType::Mcp) && t.server_url.is_some())?; - let server_url = tool.server_url.as_ref()?.trim().to_string(); - if !(server_url.starts_with("http://") || server_url.starts_with("https://")) { - warn!( - "Ignoring MCP server_url with unsupported scheme: {}", - server_url - ); - return None; - } - let name = tool - .server_label - .clone() - .unwrap_or_else(|| "request-mcp".to_string()); - let token = tool.authorization.clone(); - let transport = if server_url.contains("/sse") { - mcp::McpTransport::Sse { - url: server_url.clone(), - token, - } - } else { - mcp::McpTransport::Streamable { - url: server_url.clone(), - token, - } - }; - - // Create server config - let server_config = mcp::McpServerConfig { - name, - transport, - proxy: None, - required: false, - }; - - // Use McpManager to get or create dynamic client - match mcp_manager.get_or_create_client(server_config).await { - Ok(_client) => Some(mcp_manager.clone()), - Err(err) => { - warn!("Failed to get/create MCP connection: {}", err); - None - } - } -} - // ============================================================================ // Tool Execution // ============================================================================ @@ -705,19 +630,10 @@ pub(super) async fn execute_tool_loop( // Inject MCP output items if we executed any tools if state.total_calls > 0 { - let server_label = original_body - .tools - .as_ref() - .and_then(|tools| { - tools - .iter() - .find(|t| matches!(t.r#type, ResponseToolType::Mcp)) - .and_then(|t| t.server_label.as_deref()) - }) - .unwrap_or("mcp"); + let server_label = extract_server_label(original_body.tools.as_deref(), "mcp"); // Build mcp_list_tools item - let list_tools_item = build_mcp_list_tools_item(active_mcp, server_label); + let list_tools_item = build_mcp_list_tools_item(active_mcp, &server_label); // Insert at beginning of output array if let Some(output_array) = response_json @@ -728,7 +644,7 @@ pub(super) async fn execute_tool_loop( // Build mcp_call items using helper function let mcp_call_items = - build_executed_mcp_call_items(&state.conversation_history, server_label); + build_executed_mcp_call_items(&state.conversation_history, &server_label); // Insert mcp_call items after mcp_list_tools using mutable position let mut insert_pos = 1; @@ -767,16 +683,7 @@ pub(super) fn build_incomplete_response( // Convert any function_call in output to mcp_call format if let Some(output_array) = obj.get_mut("output").and_then(|v| v.as_array_mut()) { - let server_label = original_body - .tools - .as_ref() - .and_then(|tools| { - tools - .iter() - .find(|t| matches!(t.r#type, ResponseToolType::Mcp)) - .and_then(|t| t.server_label.as_deref()) - }) - .unwrap_or("mcp"); + let server_label = extract_server_label(original_body.tools.as_deref(), "mcp"); // Find any function_call items and convert them to mcp_call (incomplete) let mut mcp_call_items = Vec::new(); @@ -794,7 +701,7 @@ pub(super) fn build_incomplete_response( tool_name, args, "", // No output - wasn't executed - server_label, + &server_label, false, // Not successful Some("Not executed - response stopped due to limit"), ); @@ -804,12 +711,12 @@ pub(super) fn build_incomplete_response( // Add mcp_list_tools and executed mcp_call items at the beginning if state.total_calls > 0 || !mcp_call_items.is_empty() { - let list_tools_item = build_mcp_list_tools_item(active_mcp, server_label); + let list_tools_item = build_mcp_list_tools_item(active_mcp, &server_label); output_array.insert(0, list_tools_item); // Add mcp_call items for executed calls using helper let executed_items = - build_executed_mcp_call_items(&state.conversation_history, server_label); + build_executed_mcp_call_items(&state.conversation_history, &server_label); let mut insert_pos = 1; for item in executed_items { @@ -849,7 +756,7 @@ pub(super) fn build_incomplete_response( // Output Item Builders // ============================================================================ -/// Build an mcp_list_tools output item +/// Build a mcp_list_tools output item pub(super) fn build_mcp_list_tools_item(mcp: &Arc, server_label: &str) -> Value { let tools = mcp.list_tools(); let tools_json: Vec = tools @@ -874,7 +781,7 @@ pub(super) fn build_mcp_list_tools_item(mcp: &Arc, server_label }) } -/// Build an mcp_call output item +/// Build a mcp_call output item pub(super) fn build_mcp_call_item( tool_name: &str, arguments: &str, diff --git a/sgl-model-gateway/src/routers/openai/mod.rs b/sgl-model-gateway/src/routers/openai/mod.rs index 5d8d02443..02f270751 100644 --- a/sgl-model-gateway/src/routers/openai/mod.rs +++ b/sgl-model-gateway/src/routers/openai/mod.rs @@ -9,7 +9,6 @@ mod accumulator; mod context; -pub mod conversations; pub mod mcp; pub mod provider; mod responses; diff --git a/sgl-model-gateway/src/routers/openai/responses.rs b/sgl-model-gateway/src/routers/openai/responses.rs index 3b8cb1c76..8c6ee1fac 100644 --- a/sgl-model-gateway/src/routers/openai/responses.rs +++ b/sgl-model-gateway/src/routers/openai/responses.rs @@ -1,19 +1,11 @@ use serde_json::{json, Map, Value}; use tracing::warn; -use crate::{ - data_connector::{ResponseId, StoredResponse}, - protocols::{ - event_types::is_response_event, - responses::{ResponseToolType, ResponsesRequest}, - }, +use crate::protocols::{ + event_types::is_response_event, + responses::{ResponseToolType, ResponsesRequest}, }; -/// Extract a string field from JSON, returning owned String -fn get_string(json: &Value, key: &str) -> Option { - json.get(key).and_then(|v| v.as_str()).map(String::from) -} - /// Check if a JSON value is missing, null, or an empty string fn is_missing_or_empty(value: Option<&Value>) -> bool { match value { @@ -32,48 +24,6 @@ where } } -/// Build a StoredResponse from response JSON and original request -pub(super) fn build_stored_response( - response_json: &Value, - original_body: &ResponsesRequest, -) -> StoredResponse { - let mut stored = StoredResponse::new(None); - - // Initialize empty arrays - will be populated by persist_items_with_storages - stored.input = Value::Array(vec![]); - stored.output = Value::Array(vec![]); - - stored.instructions = - get_string(response_json, "instructions").or_else(|| original_body.instructions.clone()); - - stored.model = get_string(response_json, "model").or_else(|| Some(original_body.model.clone())); - - stored.safety_identifier = original_body.user.clone(); - stored.conversation_id = original_body.conversation.clone(); - - stored.metadata = response_json - .get("metadata") - .and_then(|v| v.as_object()) - .map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect()) - .unwrap_or_else(|| original_body.metadata.clone().unwrap_or_default()); - - stored.previous_response_id = get_string(response_json, "previous_response_id") - .map(|s| ResponseId::from(s.as_str())) - .or_else(|| { - original_body - .previous_response_id - .as_deref() - .map(ResponseId::from) - }); - - if let Some(id_str) = get_string(response_json, "id") { - stored.id = ResponseId::from(id_str.as_str()); - } - - stored.raw_response = response_json.clone(); - stored -} - /// Patch streaming response JSON with metadata from original request pub(super) fn patch_streaming_response_json( response_json: &mut Value, diff --git a/sgl-model-gateway/src/routers/openai/router.rs b/sgl-model-gateway/src/routers/openai/router.rs index ef7e9cec6..e62765a51 100644 --- a/sgl-model-gateway/src/routers/openai/router.rs +++ b/sgl-model-gateway/src/routers/openai/router.rs @@ -23,11 +23,7 @@ use super::{ ComponentRefs, PayloadState, RequestContext, ResponsesComponents, SharedComponents, WorkerSelection, }, - conversations::persist_conversation_items, - mcp::{ - ensure_request_mcp_client, execute_tool_loop, prepare_mcp_payload_for_streaming, - McpLoopConfig, - }, + mcp::{execute_tool_loop, prepare_mcp_payload_for_streaming}, provider::ProviderRegistry, responses::{mask_tools_as_mcp, patch_streaming_response_json}, streaming::handle_streaming_response, @@ -48,7 +44,11 @@ use crate::{ ResponsesGetParams, ResponsesRequest, }, }, - routers::header_utils::{apply_provider_headers, extract_auth_header}, + routers::{ + header_utils::{apply_provider_headers, extract_auth_header}, + mcp_utils::{ensure_request_mcp_client, McpLoopConfig}, + persistence_utils::persist_conversation_items, + }, }; pub struct OpenAIRouter { diff --git a/sgl-model-gateway/src/routers/openai/streaming.rs b/sgl-model-gateway/src/routers/openai/streaming.rs index d1e174f37..380c9c265 100644 --- a/sgl-model-gateway/src/routers/openai/streaming.rs +++ b/sgl-model-gateway/src/routers/openai/streaming.rs @@ -25,11 +25,9 @@ use tracing::warn; use super::accumulator::StreamingResponseAccumulator; use super::{ context::{RequestContext, StreamingEventContext, StreamingRequest}, - conversations::persist_conversation_items, mcp::{ - build_resume_payload, ensure_request_mcp_client, execute_streaming_tool_calls, - inject_mcp_metadata_streaming, prepare_mcp_payload_for_streaming, - send_mcp_list_tools_events, McpLoopConfig, ToolLoopState, + build_resume_payload, execute_streaming_tool_calls, inject_mcp_metadata_streaming, + prepare_mcp_payload_for_streaming, send_mcp_list_tools_events, ToolLoopState, }, responses::{mask_tools_as_mcp, patch_streaming_response_json, rewrite_streaming_block}, tool_handler::{StreamAction, StreamingToolHandler}, @@ -42,7 +40,11 @@ use crate::{ }, responses::{ResponseToolType, ResponsesRequest}, }, - routers::header_utils::{apply_request_headers, preserve_response_headers}, + routers::{ + header_utils::{apply_request_headers, preserve_response_headers}, + mcp_utils::{ensure_request_mcp_client, McpLoopConfig}, + persistence_utils::persist_conversation_items, + }, }; // ============================================================================ diff --git a/sgl-model-gateway/src/routers/openai/conversations.rs b/sgl-model-gateway/src/routers/persistence_utils.rs similarity index 60% rename from sgl-model-gateway/src/routers/openai/conversations.rs rename to sgl-model-gateway/src/routers/persistence_utils.rs index 293ed6cf2..b3775728b 100644 --- a/sgl-model-gateway/src/routers/openai/conversations.rs +++ b/sgl-model-gateway/src/routers/persistence_utils.rs @@ -1,114 +1,179 @@ -//! Conversation operations for OpenAI router -//! -//! Re-exports shared CRUD handlers and provides OpenAI-specific persistence logic. +//! Utilities for persisting responses and conversation items across router implementations. use std::sync::Arc; +use chrono::Utc; use serde_json::{json, Value}; -use tracing::{info, warn}; +use tracing::{debug, info, warn}; -use super::responses::build_stored_response; -// Re-export shared conversation handlers for backward compatibility -pub use crate::routers::conversations::{ - conversation_to_json, create_and_link_item, create_conversation, create_conversation_items, - delete_conversation, delete_conversation_item, get_conversation, get_conversation_item, - item_to_json, list_conversation_items, update_conversation, MAX_METADATA_PROPERTIES, -}; use crate::{ data_connector::{ - ConversationId, ConversationItemId, ConversationItemStorage, ConversationStorage, - NewConversationItem, ResponseId, ResponseStorage, + ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, + ConversationStorage, NewConversationItem, ResponseId, ResponseStorage, StoredResponse, }, protocols::responses::{ generate_id, ResponseInput, ResponseInputOutputItem, ResponsesRequest, StringOrContentParts, }, }; -/// Persist conversation items to storage -/// -/// This function: -/// 1. Extracts and normalizes input items from the request -/// 2. Extracts output items from the response -/// 3. Stores ALL items in response storage (always) -/// 4. If conversation provided, also links items to conversation -pub async fn persist_conversation_items( - conversation_storage: Arc, - item_storage: Arc, - response_storage: Arc, - response_json: &Value, - original_body: &ResponsesRequest, -) -> Result<(), String> { - // Extract response ID - let response_id_str = response_json - .get("id") - .and_then(|v| v.as_str()) - .ok_or_else(|| "Response missing id field".to_string())?; - let response_id = ResponseId::from(response_id_str); - // Parse and normalize input items from request - let input_items = extract_input_items(&original_body.input)?; +// ============================================================================ +// Constants +// ============================================================================ - // Parse output items from response - let output_items = response_json - .get("output") - .and_then(|v| v.as_array()) - .cloned() - .ok_or_else(|| "No output array in response".to_string())?; +/// Field mappings for item types that store data in content +pub const ITEM_TYPE_FIELDS: &[(&str, &[&str])] = &[ + ( + "mcp_call", + &[ + "name", + "arguments", + "output", + "server_label", + "approval_request_id", + "error", + ], + ), + ("mcp_list_tools", &["tools", "server_label"]), + ("function_call", &["call_id", "name", "arguments", "output"]), + ("function_call_output", &["call_id", "output"]), +]; - // Build and store response - let mut stored_response = build_stored_response(response_json, original_body); - stored_response.id = response_id.clone(); - stored_response.input = Value::Array(input_items.clone()); - stored_response.output = Value::Array(output_items.clone()); +// ============================================================================ +// JSON Serialization +// ============================================================================ - response_storage - .store_response(stored_response) - .await - .map_err(|e| format!("Failed to store response: {}", e))?; +/// Convert a ConversationItem to JSON, extracting specified fields based on item type +/// or including content as-is for standard message types. +pub fn item_to_json(item: &ConversationItem) -> Value { + let mut obj = serde_json::Map::new(); + obj.insert("id".to_string(), json!(item.id.0)); + obj.insert("type".to_string(), json!(item.item_type)); - // Check if conversation is provided and validate it exists - let conv_id_opt = if let Some(id) = &original_body.conversation { - let conv_id = ConversationId::from(id.as_str()); - match conversation_storage.get_conversation(&conv_id).await { - Ok(Some(_)) => Some(conv_id), - Ok(None) => { - warn!(conversation_id = %conv_id.0, "Conversation not found, skipping item linking"); - None + if let Some(role) = &item.role { + obj.insert("role".to_string(), json!(role)); + } + + // Find field mappings for this item type + let fields = ITEM_TYPE_FIELDS + .iter() + .find(|(t, _)| *t == item.item_type) + .map(|(_, fields)| *fields); + + if let Some(fields) = fields { + // Extract specific fields from content + if let Some(content_obj) = item.content.as_object() { + for field in fields { + if let Some(value) = content_obj.get(*field) { + obj.insert((*field).to_string(), value.clone()); + } } - Err(e) => return Err(format!("Failed to get conversation: {}", e)), } } else { - None - }; + // Default: include content as-is + obj.insert("content".to_string(), item.content.clone()); + } + + if let Some(status) = &item.status { + obj.insert("status".to_string(), json!(status)); + } + + Value::Object(obj) +} + +// ============================================================================ +// Item Creation Helper +// ============================================================================ + +/// Create a conversation item and optionally link it to a conversation. +/// Sets default "completed" status if not provided. +pub async fn create_and_link_item( + item_storage: &Arc, + conv_id_opt: Option<&ConversationId>, + mut new_item: NewConversationItem, +) -> Result<(), String> { + if new_item.status.is_none() { + new_item.status = Some("completed".to_string()); + } + + let created = item_storage + .create_item(new_item) + .await + .map_err(|e| format!("Failed to create item: {e}"))?; - // If conversation exists, link items to it if let Some(conv_id) = conv_id_opt { - link_items_to_conversation( - &item_storage, - &conv_id, - &input_items, - &output_items, - response_id_str, - ) - .await?; - info!( + item_storage + .link_item(conv_id, &created.id, Utc::now()) + .await + .map_err(|e| format!("Failed to link item: {e}"))?; + + debug!( conversation_id = %conv_id.0, - response_id = %response_id.0, - input_count = input_items.len(), - output_count = output_items.len(), - "Persisted response and linked items to conversation" + item_id = %created.id.0, + item_type = %created.item_type, + "Persisted conversation item and link" ); } else { - info!( - response_id = %response_id.0, - input_count = input_items.len(), - output_count = output_items.len(), - "Persisted response without conversation linking" + debug!( + item_id = %created.id.0, + item_type = %created.item_type, + "Persisted conversation item (no conversation link)" ); } Ok(()) } +// ============================================================================ +// Response Persistence +// ============================================================================ + +/// Extract a string field from JSON, returning owned String +fn get_string(json: &Value, key: &str) -> Option { + json.get(key).and_then(|v| v.as_str()).map(String::from) +} + +/// Build a StoredResponse from response JSON and original request +pub fn build_stored_response( + response_json: &Value, + original_body: &ResponsesRequest, +) -> StoredResponse { + let mut stored = StoredResponse::new(None); + + // Initialize empty arrays - will be populated by persist_conversation_items + stored.input = Value::Array(vec![]); + stored.output = Value::Array(vec![]); + + stored.instructions = + get_string(response_json, "instructions").or_else(|| original_body.instructions.clone()); + + stored.model = get_string(response_json, "model").or_else(|| Some(original_body.model.clone())); + + stored.safety_identifier = original_body.user.clone(); + stored.conversation_id = original_body.conversation.clone(); + + stored.metadata = response_json + .get("metadata") + .and_then(|v| v.as_object()) + .map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect()) + .unwrap_or_else(|| original_body.metadata.clone().unwrap_or_default()); + + stored.previous_response_id = get_string(response_json, "previous_response_id") + .map(|s| ResponseId::from(s.as_str())) + .or_else(|| { + original_body + .previous_response_id + .as_deref() + .map(ResponseId::from) + }); + + if let Some(id_str) = get_string(response_json, "id") { + stored.id = ResponseId::from(id_str.as_str()); + } + + stored.raw_response = response_json.clone(); + stored +} + /// Extract and normalize input items from ResponseInput fn extract_input_items(input: &ResponseInput) -> Result, String> { let items = match input { @@ -149,7 +214,7 @@ fn extract_input_items(input: &ResponseInput) -> Result, String> { })) } _ => { - // For other item types (Message, Reasoning, FunctionToolCall, FunctionCallOutput), serialize and ensure ID + // For other item types, serialize and ensure ID let mut value = serde_json::to_value(item) .map_err(|e| format!("Failed to serialize item: {}", e))?; @@ -253,3 +318,89 @@ async fn link_items_to_conversation( Ok(()) } + +/// Persist conversation items to storage +/// +/// This function: +/// 1. Extracts and normalizes input items from the request +/// 2. Extracts output items from the response +/// 3. Stores ALL items in response storage (always) +/// 4. If conversation provided, also links items to conversation +pub async fn persist_conversation_items( + conversation_storage: Arc, + item_storage: Arc, + response_storage: Arc, + response_json: &Value, + original_body: &ResponsesRequest, +) -> Result<(), String> { + // Extract response ID + let response_id_str = response_json + .get("id") + .and_then(|v| v.as_str()) + .ok_or_else(|| "Response missing id field".to_string())?; + let response_id = ResponseId::from(response_id_str); + + // Parse and normalize input items from request + let input_items = extract_input_items(&original_body.input)?; + + // Parse output items from response + let output_items = response_json + .get("output") + .and_then(|v| v.as_array()) + .cloned() + .ok_or_else(|| "No output array in response".to_string())?; + + // Build and store response + let mut stored_response = build_stored_response(response_json, original_body); + stored_response.id = response_id.clone(); + stored_response.input = Value::Array(input_items.clone()); + stored_response.output = Value::Array(output_items.clone()); + + response_storage + .store_response(stored_response) + .await + .map_err(|e| format!("Failed to store response: {}", e))?; + + // Check if conversation is provided and validate it exists + let conv_id_opt = if let Some(id) = &original_body.conversation { + let conv_id = ConversationId::from(id.as_str()); + match conversation_storage.get_conversation(&conv_id).await { + Ok(Some(_)) => Some(conv_id), + Ok(None) => { + warn!(conversation_id = %conv_id.0, "Conversation not found, skipping item linking"); + None + } + Err(e) => return Err(format!("Failed to get conversation: {}", e)), + } + } else { + None + }; + + // If conversation exists, link items to it + if let Some(conv_id) = conv_id_opt { + link_items_to_conversation( + &item_storage, + &conv_id, + &input_items, + &output_items, + response_id_str, + ) + .await?; + info!( + conversation_id = %conv_id.0, + response_id = %response_id.0, + input_count = input_items.len(), + output_count = output_items.len(), + "Persisted response and linked items to conversation" + ); + } else { + info!( + response_id = %response_id.0, + input_count = input_items.len(), + output_count = output_items.len(), + "Persisted response without conversation linking" + ); + } + + Ok(()) +}