chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[model-gateway] Clear architectual debt in responses API (#16359)

Browse Source
This commit is contained in:
Chang Su
2026-01-03 21:35:28 -08:00
committed by GitHub
parent 5f3eb377e0
commit 8d58b3dcdd
13 changed files with 441 additions and 421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/routers/conversations/handlers.rs
+9 -100
View File
@@ -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<dyn ConversationItemStorage>,
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,
sgl-model-gateway/src/routers/grpc/common/responses/utils.rs
+1 -2
View File
@@ -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,
},
};
sgl-model-gateway/src/routers/grpc/harmony/responses.rs
+18 -32
View File
@@ -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(&current_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<Result<Bytes, std::io::Error>>,
) {
// 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(&current_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
sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs
+13 -33
View File
@@ -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<String>,
) -> Result<ResponsesResponse, Response> {
// 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<Bytes, std::io::Error>>,
) -> 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;
}
sgl-model-gateway/src/routers/grpc/router.rs
-1
View File
@@ -37,7 +37,6 @@ use crate::{
/// gRPC router implementation for SGLang
#[derive(Clone)]
#[allow(dead_code)]
pub struct GrpcRouter {
worker_registry: Arc<WorkerRegistry>,
pipeline: RequestPipeline,
sgl-model-gateway/src/routers/mcp_utils.rs
+136
View File
@@ -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<McpManager>,
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
}
}
}
sgl-model-gateway/src/routers/mod.rs
+2
View File
@@ -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;
sgl-model-gateway/src/routers/openai/mcp.rs
+14 -107
View File
@@ -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<mcp::McpManager>,
tools: &[ResponseTool],
) -> Option<Arc<mcp::McpManager>> {
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<mcp::McpManager>, server_label: &str) -> Value {
let tools = mcp.list_tools();
let tools_json: Vec<Value> = tools
@@ -874,7 +781,7 @@ pub(super) fn build_mcp_list_tools_item(mcp: &Arc<mcp::McpManager>, 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,
sgl-model-gateway/src/routers/openai/mod.rs
-1
View File
@@ -9,7 +9,6 @@
mod accumulator;
mod context;
pub mod conversations;
pub mod mcp;
pub mod provider;
mod responses;
sgl-model-gateway/src/routers/openai/responses.rs
+3 -53
View File
@@ -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<String> {
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,
sgl-model-gateway/src/routers/openai/router.rs
+6 -6
View File
@@ -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 {
sgl-model-gateway/src/routers/openai/streaming.rs
+7 -5
View File
@@ -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,
},
};
// ============================================================================
sgl-model-gateway/src/routers/openai/conversations.rs → sgl-model-gateway/src/routers/persistence_utils.rs
+232 -81
View File
@@ -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<dyn ConversationStorage>,
item_storage: Arc<dyn ConversationItemStorage>,
response_storage: Arc<dyn ResponseStorage>,
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<dyn ConversationItemStorage>,
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<String> {
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<Vec<Value>, String> {
let items = match input {
@@ -149,7 +214,7 @@ fn extract_input_items(input: &ResponseInput) -> Result<Vec<Value>, 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<dyn ConversationStorage>,
item_storage: Arc<dyn ConversationItemStorage>,
response_storage: Arc<dyn ResponseStorage>,
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(())
}