[router] add basic ci tests for gpt-oss model support (#12651)

Co-authored-by: Chang Su <chang.s.su@oracle.com>
This commit is contained in:
Keyang Ru
2025-11-05 15:45:11 -08:00
committed by GitHub
co-authored by Chang Su
parent b0476a0642
commit 7ef1964a66
21 changed files with 1008 additions and 241 deletions
+2
View File
@@ -128,6 +128,7 @@ pub enum ResponseInputOutputItem {
id: String,
summary: Vec<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
#[serde(default)]
content: Vec<ResponseReasoningContent>,
#[serde(skip_serializing_if = "Option::is_none")]
status: Option<String>,
@@ -168,6 +169,7 @@ pub enum ResponseContentPart {
#[serde(rename = "output_text")]
OutputText {
text: String,
#[serde(default)]
#[serde(skip_serializing_if = "Vec::is_empty")]
annotations: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
@@ -6,4 +6,4 @@ pub mod utils;
pub use handlers::{cancel_response_impl, get_response_impl};
pub use streaming::{build_sse_response, OutputItemType, ResponseStreamEventEmitter};
pub use utils::ensure_mcp_connection;
pub use utils::{ensure_mcp_connection, persist_response_if_needed};
@@ -11,7 +11,14 @@ use uuid::Uuid;
use crate::{
mcp,
protocols::{chat::ChatCompletionStreamResponse, responses::ResponsesRequest},
protocols::{
chat::ChatCompletionStreamResponse,
common::{Usage, UsageInfo},
responses::{
ResponseOutputItem, ResponseStatus, ResponsesRequest, ResponsesResponse, ResponsesUsage,
},
},
routers::grpc::harmony::responses::ToolResult,
};
pub enum OutputItemType {
@@ -69,6 +76,7 @@ pub struct ResponseStreamEventEmitter {
has_emitted_content_part_added: bool,
// MCP call tracking
mcp_call_accumulated_args: HashMap<String, String>,
pub(crate) mcp_server_label: Option<String>, // Server label for MCP tools
// Output item tracking
output_items: Vec<OutputItemState>,
next_output_index: usize,
@@ -93,6 +101,7 @@ impl ResponseStreamEventEmitter {
has_emitted_output_item_added: false,
has_emitted_content_part_added: false,
mcp_call_accumulated_args: HashMap::new(),
mcp_server_label: None,
output_items: Vec::new(),
next_output_index: 0,
current_message_output_index: None,
@@ -106,6 +115,41 @@ impl ResponseStreamEventEmitter {
self.original_request = Some(request);
}
/// Set the MCP server label for MCP tool calls
pub fn set_mcp_server_label(&mut self, server_label: String) {
self.mcp_server_label = Some(server_label);
}
/// Update mcp_call output items with tool execution results
///
/// After MCP tools are executed, this updates the stored output items
/// to include the output field from the tool results.
pub(crate) fn update_mcp_call_outputs(&mut self, tool_results: &[ToolResult]) {
for tool_result in tool_results {
// Find the output item with matching call_id
for item_state in self.output_items.iter_mut() {
if let Some(ref mut item_data) = item_state.item_data {
// Check if this is an mcp_call item with matching call_id
if item_data.get("type").and_then(|t| t.as_str()) == Some("mcp_call")
&& item_data.get("call_id").and_then(|c| c.as_str())
== Some(&tool_result.call_id)
{
// Add output field
let output_str = serde_json::to_string(&tool_result.output)
.unwrap_or_else(|_| "{}".to_string());
item_data["output"] = json!(output_str);
// Update status based on success
if tool_result.is_error {
item_data["status"] = json!("failed");
}
break;
}
}
}
}
}
fn next_sequence(&mut self) -> u64 {
let seq = self.sequence_number;
self.sequence_number += 1;
@@ -542,6 +586,78 @@ impl ResponseStreamEventEmitter {
}
}
/// Finalize and return the complete ResponsesResponse
///
/// This constructs the final ResponsesResponse from all accumulated output items
/// for persistence. Should be called after streaming is complete.
pub fn finalize(&self, usage: Option<Usage>) -> ResponsesResponse {
// Build output array from tracked items
let output: Vec<ResponseOutputItem> = self
.output_items
.iter()
.filter_map(|item| {
item.item_data
.as_ref()
.and_then(|data| serde_json::from_value(data.clone()).ok())
})
.collect();
// Convert Usage to ResponsesUsage
let responses_usage = usage.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)
});
// Get original request fields or use defaults
let req = self.original_request.as_ref();
// Convert tool_choice to String
let tool_choice = req
.and_then(|r| r.tool_choice.as_ref())
.map(|tc| serde_json::to_string(tc).unwrap_or_else(|_| "auto".to_string()))
.unwrap_or_else(|| "auto".to_string());
ResponsesResponse {
id: self.response_id.clone(),
object: "response".to_string(),
created_at: self.created_at as i64,
status: ResponseStatus::Completed,
error: None,
incomplete_details: None,
instructions: req.and_then(|r| r.instructions.clone()),
max_output_tokens: req.and_then(|r| r.max_output_tokens),
model: self.model.clone(),
output,
parallel_tool_calls: req.and_then(|r| r.parallel_tool_calls).unwrap_or(true),
previous_response_id: req.and_then(|r| r.previous_response_id.clone()),
reasoning: None, // TODO: Extract from output items if needed
store: req.and_then(|r| r.store).unwrap_or(true),
temperature: req.and_then(|r| r.temperature),
text: None,
tool_choice,
tools: req
.map(|r| r.tools.clone().unwrap_or_default())
.unwrap_or_default(),
top_p: req.and_then(|r| r.top_p),
truncation: None, // Convert from Truncation to String if needed
usage: responses_usage,
metadata: req
.map(|r| r.metadata.clone().unwrap_or_default())
.unwrap_or_default(),
user: req.and_then(|r| r.user.clone()),
safety_identifier: None,
}
}
/// Emit reasoning item wrapper events (added + done)
///
/// Reasoning items in OpenAI format are simple placeholders emitted between tool iterations.
@@ -6,16 +6,21 @@ use axum::{
http::StatusCode,
response::{IntoResponse, Response},
};
use serde_json::json;
use serde_json::{json, to_value};
use tracing::{debug, warn};
use crate::{
core::WorkerRegistry,
data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage},
mcp::McpManager,
protocols::{
common::Tool,
responses::{ResponseTool, ResponseToolType},
responses::{ResponseTool, ResponseToolType, ResponsesRequest, ResponsesResponse},
},
routers::{
grpc::error,
openai::{conversations::persist_conversation_items, mcp::ensure_request_mcp_client},
},
routers::{grpc::error, openai::mcp::ensure_request_mcp_client},
};
/// Ensure MCP connection succeeds if MCP tools are declared
@@ -123,3 +128,35 @@ pub fn extract_tools_from_response_tools(
})
.collect()
}
/// Persist response to storage if store=true
///
/// Common helper function to avoid duplication across sync and streaming paths
/// in both harmony and regular responses implementations.
pub async fn persist_response_if_needed(
conversation_storage: Arc<dyn ConversationStorage>,
conversation_item_storage: Arc<dyn ConversationItemStorage>,
response_storage: Arc<dyn ResponseStorage>,
response: &ResponsesResponse,
original_request: &ResponsesRequest,
) {
if !original_request.store.unwrap_or(true) {
return;
}
if let Ok(response_json) = to_value(response) {
if let Err(e) = persist_conversation_items(
conversation_storage,
conversation_item_storage,
response_storage,
&response_json,
original_request,
)
.await
{
warn!("Failed to persist response: {}", e);
} else {
debug!("Persisted response: {}", response.id);
}
}
}
+121 -26
View File
@@ -49,7 +49,7 @@ use tracing::{debug, warn};
use uuid::Uuid;
use crate::{
data_connector::{ResponseId, ResponseStorage},
data_connector::{ConversationItemStorage, ConversationStorage, ResponseId, ResponseStorage},
mcp::{self, McpManager},
protocols::{
common::{Function, ToolCall, ToolChoice, ToolChoiceValue, Usage},
@@ -62,7 +62,7 @@ use crate::{
},
routers::grpc::{
common::responses::{
build_sse_response, ensure_mcp_connection,
build_sse_response, ensure_mcp_connection, persist_response_if_needed,
streaming::{OutputItemType, ResponseStreamEventEmitter},
},
context::SharedComponents,
@@ -157,6 +157,12 @@ pub struct HarmonyResponsesContext {
/// Response storage for loading conversation history
pub response_storage: Arc<dyn ResponseStorage>,
/// Conversation storage for persisting conversations
pub conversation_storage: Arc<dyn ConversationStorage>,
/// Conversation item storage for persisting conversation items
pub conversation_item_storage: Arc<dyn ConversationItemStorage>,
/// Optional streaming sender (for future streaming support)
pub stream_tx: Option<mpsc::UnboundedSender<Result<String, String>>>,
}
@@ -168,12 +174,16 @@ impl HarmonyResponsesContext {
components: Arc<SharedComponents>,
mcp_manager: Arc<McpManager>,
response_storage: Arc<dyn ResponseStorage>,
conversation_storage: Arc<dyn ConversationStorage>,
conversation_item_storage: Arc<dyn ConversationItemStorage>,
) -> Self {
Self {
pipeline,
components,
mcp_manager,
response_storage,
conversation_storage,
conversation_item_storage,
stream_tx: None,
}
}
@@ -184,6 +194,8 @@ impl HarmonyResponsesContext {
components: Arc<SharedComponents>,
mcp_manager: Arc<McpManager>,
response_storage: Arc<dyn ResponseStorage>,
conversation_storage: Arc<dyn ConversationStorage>,
conversation_item_storage: Arc<dyn ConversationItemStorage>,
stream_tx: mpsc::UnboundedSender<Result<String, String>>,
) -> Self {
Self {
@@ -191,6 +203,8 @@ impl HarmonyResponsesContext {
components,
mcp_manager,
response_storage,
conversation_storage,
conversation_item_storage,
stream_tx: Some(stream_tx),
}
}
@@ -237,6 +251,9 @@ pub async fn serve_harmony_responses(
ctx: &HarmonyResponsesContext,
request: ResponsesRequest,
) -> Result<ResponsesResponse, Response> {
// Clone request for persistence
let original_request = request.clone();
// Load previous conversation history if previous_response_id is set
let current_request = load_previous_messages(ctx, request).await?;
@@ -244,12 +261,24 @@ pub async fn serve_harmony_responses(
let has_mcp_tools =
ensure_mcp_connection(&ctx.mcp_manager, current_request.tools.as_deref()).await?;
if has_mcp_tools {
execute_with_mcp_loop(ctx, current_request).await
let response = if has_mcp_tools {
execute_with_mcp_loop(ctx, current_request).await?
} else {
// No MCP tools - execute pipeline once (may have function tools or no tools)
execute_without_mcp_loop(ctx, current_request).await
}
execute_without_mcp_loop(ctx, current_request).await?
};
// Persist response to storage if store=true
persist_response_if_needed(
ctx.conversation_storage.clone(),
ctx.conversation_item_storage.clone(),
ctx.response_storage.clone(),
&response,
&original_request,
)
.await;
Ok(response)
}
/// Execute Harmony Responses with MCP tool loop
@@ -260,7 +289,10 @@ async fn execute_with_mcp_loop(
mut current_request: ResponsesRequest,
) -> Result<ResponsesResponse, Response> {
let mut iteration_count = 0;
let mut mcp_tracking = McpCallTracking::new("sglang-mcp".to_string());
// Extract server_label from request tools
let server_label = extract_mcp_server_label(current_request.tools.as_deref());
let mut mcp_tracking = McpCallTracking::new(server_label.clone());
// Extract user's max_tool_calls limit (if set)
let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
@@ -472,7 +504,7 @@ pub async fn serve_harmony_responses_stream(
request: ResponsesRequest,
) -> Response {
// Load previous conversation history if previous_response_id is set
let current_request = match load_previous_messages(ctx, request).await {
let current_request = match load_previous_messages(ctx, request.clone()).await {
Ok(req) => req,
Err(err_response) => return err_response,
};
@@ -517,9 +549,10 @@ pub async fn serve_harmony_responses_stream(
}
if has_mcp_tools {
execute_mcp_tool_loop_streaming(ctx, current_request, &mut emitter, &tx).await;
execute_mcp_tool_loop_streaming(ctx, current_request, &request, &mut emitter, &tx)
.await;
} else {
execute_without_mcp_streaming(ctx, &current_request, &mut emitter, &tx).await;
execute_without_mcp_streaming(ctx, &current_request, &request, &mut emitter, &tx).await;
}
});
@@ -534,14 +567,22 @@ pub async fn serve_harmony_responses_stream(
/// - Emits mcp_list_tools events
/// - Loops through tool execution iterations
/// - Emits final response.completed event
/// - Persists response internally
async fn execute_mcp_tool_loop_streaming(
ctx: &HarmonyResponsesContext,
mut current_request: ResponsesRequest,
original_request: &ResponsesRequest,
emitter: &mut ResponseStreamEventEmitter,
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());
// Set server label in emitter for MCP call items
emitter.set_mcp_server_label(server_label.clone());
// Initialize MCP call tracking
let mut mcp_tracking = McpCallTracking::new("sglang-mcp".to_string());
let mut mcp_tracking = McpCallTracking::new(server_label.clone());
// Extract user's max_tool_calls limit (if set)
let max_tool_calls = current_request.max_tool_calls.map(|n| n as usize);
@@ -576,11 +617,26 @@ async fn execute_mcp_tool_loop_streaming(
})
.collect();
// Build final item with completed status and tools
let item_done = json!({
"id": item_id,
"type": "mcp_list_tools",
"server_label": server_label,
"status": "completed",
"tools": tool_items
});
// Store the completed item data and mark as completed FIRST
// This ensures it appears in final response even if event sending fails
emitter.emit_output_item_done(output_index, &item_done);
emitter.complete_output_item(output_index);
// Now emit all the events (failures won't affect the stored data)
// Emit output_item.added
let item = json!({
"id": item_id,
"type": "mcp_list_tools",
"server_label": "sglang-mcp",
"server_label": server_label,
"status": "in_progress",
"tools": []
});
@@ -602,20 +658,11 @@ async fn execute_mcp_tool_loop_streaming(
}
// Emit output_item.done
let item_done = json!({
"id": item_id,
"type": "mcp_list_tools",
"server_label": "sglang-mcp",
"status": "completed",
"tools": tool_items
});
let event = emitter.emit_output_item_done(output_index, &item_done);
if emitter.send_event(&event, tx).is_err() {
return;
}
emitter.complete_output_item(output_index);
debug!(
tool_count = mcp_tools.len(),
"Emitted mcp_list_tools on first iteration"
@@ -736,6 +783,9 @@ async fn execute_mcp_tool_loop_streaming(
}
};
// Update mcp_call output items with execution results
emitter.update_mcp_call_outputs(&tool_results);
// Build next request with appended history
current_request = match build_next_request_with_tools(
current_request,
@@ -765,6 +815,19 @@ async fn execute_mcp_tool_loop_streaming(
"Harmony Responses streaming completed - no more tool calls"
);
// Finalize response from emitter's accumulated data
let final_response = emitter.finalize(Some(usage.clone()));
// Persist response to storage if store=true
persist_response_if_needed(
ctx.conversation_storage.clone(),
ctx.conversation_item_storage.clone(),
ctx.response_storage.clone(),
&final_response,
original_request,
)
.await;
// Emit response.completed with usage
let usage_json = json!({
"input_tokens": usage.prompt_tokens,
@@ -786,6 +849,7 @@ async fn execute_mcp_tool_loop_streaming(
async fn execute_without_mcp_streaming(
ctx: &HarmonyResponsesContext,
current_request: &ResponsesRequest,
original_request: &ResponsesRequest,
emitter: &mut ResponseStreamEventEmitter,
tx: &mpsc::UnboundedSender<Result<Bytes, std::io::Error>>,
) {
@@ -830,6 +894,19 @@ async fn execute_without_mcp_streaming(
ResponsesIterationResult::Completed { usage, .. } => usage,
};
// Finalize response from emitter's accumulated data
let final_response = emitter.finalize(Some(usage.clone()));
// Persist response to storage if store=true
persist_response_if_needed(
ctx.conversation_storage.clone(),
ctx.conversation_item_storage.clone(),
ctx.response_storage.clone(),
&final_response,
original_request,
)
.await;
// Emit response.completed with usage
let usage_json = json!({
"input_tokens": usage.prompt_tokens,
@@ -1191,19 +1268,19 @@ fn build_next_request_with_tools(
/// Tool execution result
///
/// Contains the result of executing a single MCP tool.
struct ToolResult {
pub(crate) struct ToolResult {
/// Tool call ID (for matching with request)
call_id: String,
pub(crate) call_id: String,
/// Tool name
#[allow(dead_code)] // Kept for documentation and future use
tool_name: String,
pub(crate) tool_name: String,
/// Tool output (JSON value)
output: Value,
pub(crate) output: Value,
/// Whether this is an error result
is_error: bool,
pub(crate) is_error: bool,
}
/// Convert MCP tools to Responses API tool format
@@ -1297,6 +1374,24 @@ 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
@@ -733,7 +733,7 @@ impl HarmonyStreamingProcessor {
// Emit output_item.added wrapper event
let call_id = tc_delta.id.as_ref().unwrap();
let item = json!({
let mut item = json!({
"id": item_id,
"type": mode.type_str(),
"name": tool_name,
@@ -741,6 +741,14 @@ impl HarmonyStreamingProcessor {
"arguments": "",
"status": "in_progress"
});
// Add server_label for MCP calls
if mode.emits_status_events() {
if let Some(ref server_label) = emitter.mcp_server_label {
item["server_label"] = json!(server_label);
}
}
let event = emitter.emit_output_item_added(output_index, &item);
emitter.send_event_best_effort(&event, tx);
@@ -836,7 +844,7 @@ impl HarmonyStreamingProcessor {
}
// Emit output_item.done wrapper event
let item = json!({
let mut item = json!({
"id": item_id,
"type": mode.type_str(),
"name": tool_name,
@@ -844,11 +852,21 @@ impl HarmonyStreamingProcessor {
"arguments": args_str,
"status": "completed"
});
let event = emitter.emit_output_item_done(*output_index, &item);
emitter.send_event_best_effort(&event, tx);
// Mark output item as completed
// Add server_label for MCP calls
if mode.emits_status_events() {
// MCP mode - include server_label
if let Some(ref server_label) = emitter.mcp_server_label {
item["server_label"] = json!(server_label);
}
}
let event = emitter.emit_output_item_done(*output_index, &item);
// Mark output item as completed before sending
emitter.complete_output_item(*output_index);
emitter.send_event_best_effort(&event, tx);
}
}
}
@@ -878,9 +896,11 @@ impl HarmonyStreamingProcessor {
}]
});
let event = emitter.emit_output_item_done(output_index, &item);
emitter.send_event_best_effort(&event, tx);
// Mark as completed before sending (so it's included in final output even if send fails)
emitter.complete_output_item(output_index);
emitter.send_event_best_effort(&event, tx);
}
}
Some(proto::generate_response::Response::Error(err)) => {
@@ -939,7 +959,7 @@ impl HarmonyStreamingProcessor {
}
// Emit output_item.done wrapper event
let item = json!({
let mut item = json!({
"id": item_id,
"type": mode.type_str(),
"name": tool_name,
@@ -947,11 +967,20 @@ impl HarmonyStreamingProcessor {
"arguments": args_str,
"status": "completed"
});
let event = emitter.emit_output_item_done(*output_index, &item);
emitter.send_event_best_effort(&event, tx);
// Mark output item as completed
// Add server_label for MCP calls
if mode.emits_status_events() {
if let Some(ref server_label) = emitter.mcp_server_label {
item["server_label"] = json!(server_label);
}
}
let event = emitter.emit_output_item_done(*output_index, &item);
// Mark output item as completed before sending
emitter.complete_output_item(*output_index);
emitter.send_event_best_effort(&event, tx);
}
}
}
@@ -65,14 +65,12 @@ use crate::{
ResponsesUsage,
},
},
routers::{
grpc::{
common::responses::{
build_sse_response, ensure_mcp_connection, streaming::ResponseStreamEventEmitter,
},
error,
routers::grpc::{
common::responses::{
build_sse_response, ensure_mcp_connection, persist_response_if_needed,
streaming::ResponseStreamEventEmitter,
},
openai::conversations::persist_conversation_items,
error,
},
};
@@ -221,21 +219,14 @@ async fn route_responses_internal(
};
// 5. Persist response to storage if store=true
if request.store.unwrap_or(true) {
if let Ok(response_json) = serde_json::to_value(&responses_response) {
if let Err(e) = persist_conversation_items(
ctx.conversation_storage.clone(),
ctx.conversation_item_storage.clone(),
ctx.response_storage.clone(),
&response_json,
&request,
)
.await
{
warn!("Failed to persist response: {}", e);
}
}
}
persist_response_if_needed(
ctx.conversation_storage.clone(),
ctx.conversation_item_storage.clone(),
ctx.response_storage.clone(),
&responses_response,
&request,
)
.await;
Ok(responses_response)
}
@@ -454,25 +445,15 @@ async fn process_and_transform_sse_stream(
event_emitter.send_event(&completed_event, &tx)?;
// Finalize and persist accumulated response
if original_request.store.unwrap_or(true) {
let final_response = accumulator.finalize();
if let Ok(response_json) = serde_json::to_value(&final_response) {
if let Err(e) = persist_conversation_items(
conversation_storage.clone(),
conversation_item_storage.clone(),
response_storage.clone(),
&response_json,
&original_request,
)
.await
{
warn!("Failed to persist streaming response: {}", e);
} else {
debug!("Persisted streaming response: {}", final_response.id);
}
}
}
let final_response = accumulator.finalize();
persist_response_if_needed(
conversation_storage,
conversation_item_storage,
response_storage,
&final_response,
&original_request,
)
.await;
Ok(())
}
+4
View File
@@ -216,6 +216,10 @@ impl GrpcRouter {
self.shared_components.clone(),
self.harmony_responses_context.mcp_manager.clone(),
self.harmony_responses_context.response_storage.clone(),
self.harmony_responses_context.conversation_storage.clone(),
self.harmony_responses_context
.conversation_item_storage
.clone(),
);
if body.stream.unwrap_or(false) {