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

[grpc] Unify ResponsesContext and HarmonyResponsesContext (#16549)

Browse Source
This commit is contained in:
Chang Su
2026-01-06 07:28:06 -08:00
committed by GitHub
parent ed307a40bf
commit 21da2dc1ce
16 changed files with 69 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
sgl-model-gateway/src/routers/grpc/harmony/responses/context.rs → sgl-model-gateway/src/routers/grpc/common/responses/context.rs
View File

@@ -1,4 +1,6 @@
//! Context for Harmony Responses execution
//! Shared context for /v1/responses endpoint handlers
//!
//! This context is used by both regular and harmony response implementations.
use std::sync::{Arc, RwLock as StdRwLock};
@@ -8,52 +10,52 @@ use crate::{
routers::grpc::{context::SharedComponents, pipeline::RequestPipeline},
};
/// Context for Harmony Responses execution with MCP tool support
/// Context for /v1/responses endpoint
///
/// Contains all dependencies needed for multi-turn Responses API execution.
/// Cheap to clone (all Arc references).
/// Used by both regular and harmony implementations.
/// All fields are Arc/shared references, so cloning this context is cheap.
#[derive(Clone)]
pub(crate) struct HarmonyResponsesContext {
/// Pipeline for executing Harmony requests
pub(crate) struct ResponsesContext {
/// Chat pipeline for executing requests
pub pipeline: Arc<RequestPipeline>,
/// Shared components (tokenizer, parsers)
pub components: Arc<SharedComponents>,
/// MCP manager for tool execution
/// Response storage backend
pub response_storage: Arc<dyn ResponseStorage>,
/// Conversation storage backend
pub conversation_storage: Arc<dyn ConversationStorage>,
/// Conversation item storage backend
pub conversation_item_storage: Arc<dyn ConversationItemStorage>,
/// MCP manager for tool support
pub mcp_manager: Arc<McpManager>,
/// Server keys for MCP tools requested in this context
pub requested_servers: Arc<StdRwLock<Vec<String>>>,
/// 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>,
}
impl HarmonyResponsesContext {
/// Create a new Harmony Responses context
impl ResponsesContext {
/// Create a new responses context
pub fn new(
pipeline: Arc<RequestPipeline>,
components: Arc<SharedComponents>,
mcp_manager: Arc<McpManager>,
response_storage: Arc<dyn ResponseStorage>,
conversation_storage: Arc<dyn ConversationStorage>,
conversation_item_storage: Arc<dyn ConversationItemStorage>,
mcp_manager: Arc<McpManager>,
) -> Self {
Self {
pipeline,
components,
mcp_manager,
requested_servers: Arc::new(StdRwLock::new(Vec::new())),
response_storage,
conversation_storage,
conversation_item_storage,
mcp_manager,
requested_servers: Arc::new(StdRwLock::new(Vec::new())),
}
}
}

90
sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs
View File

@@ -2,17 +2,10 @@
//!
//! These handlers are used by both pipelines for retrieving and cancelling responses.
use axum::{
http::StatusCode,
response::{IntoResponse, Response},
};
use serde_json::json;
use tracing::{debug, error, warn};
use axum::response::{IntoResponse, Response};
use crate::{
data_connector::ResponseId,
routers::{error, grpc::regular::responses::ResponsesContext},
};
use super::ResponsesContext;
use crate::{data_connector::ResponseId, routers::error};
/// Implementation for GET /v1/responses/{response_id}
///
@@ -37,14 +30,14 @@ pub(crate) async fn get_response_impl(ctx: &ResponsesContext, response_id: &str)
/// Implementation for POST /v1/responses/{response_id}/cancel
///
/// Cancels a background response if it's still in progress.
/// Background mode is no longer supported, so this endpoint always returns
/// an error indicating that cancellation is not available.
pub(crate) async fn cancel_response_impl(ctx: &ResponsesContext, response_id: &str) -> Response {
let resp_id = ResponseId::from(response_id);
// Retrieve response from storage to check if it exists and get current status
// Check if response exists
match ctx.response_storage.get_response(&resp_id).await {
Ok(Some(stored_response)) => {
// Check current status - only queued or in_progress responses can be cancelled
let current_status = stored_response
.raw_response
.get("status")
@@ -52,60 +45,6 @@ pub(crate) async fn cancel_response_impl(ctx: &ResponsesContext, response_id: &s
.unwrap_or("unknown");
match current_status {
"queued" | "in_progress" => {
// Attempt to abort the background task
let mut tasks = ctx.background_tasks.write().await;
if let Some(task_info) = tasks.remove(response_id) {
// Abort the Rust task immediately
task_info.handle.abort();
// Abort the Python/scheduler request via gRPC (if client is available)
let client_opt = task_info.client.read().await;
if let Some(ref client) = *client_opt {
if let Err(e) = client
.abort_request(
task_info.grpc_request_id.clone(),
"User cancelled via API".to_string(),
)
.await
{
warn!(
"Failed to abort Python request {}: {}",
task_info.grpc_request_id, e
);
} else {
debug!(
"Successfully aborted Python request: {}",
task_info.grpc_request_id
);
}
} else {
debug!("Client not yet available for abort, request may not have started yet");
}
// Task was found and aborted
(
StatusCode::OK,
axum::Json(json!({
"id": response_id,
"status": "cancelled",
"message": "Background task has been cancelled"
})),
)
.into_response()
} else {
// Task handle not found but status is queued/in_progress
// This can happen if: (1) task crashed, or (2) storage persistence failed
error!(
"Response {} has status '{}' but task handle is missing. Task may have crashed or storage update failed.",
response_id, current_status
);
error::internal_error(
"status_update_failed",
"Internal error: background task completed but failed to update status in storage",
)
}
}
"completed" => error::bad_request(
"response_already_completed",
"Cannot cancel completed response",
@@ -113,20 +52,11 @@ pub(crate) async fn cancel_response_impl(ctx: &ResponsesContext, response_id: &s
"failed" => {
error::bad_request("response_already_failed", "Cannot cancel failed response")
}
"cancelled" => (
StatusCode::OK,
axum::Json(json!({
"id": response_id,
"status": "cancelled",
"message": "Response was already cancelled"
})),
)
.into_response(),
_ => {
// Unknown status
error::internal_error(
"unknown_response_status",
format!("Unknown response status: {}", current_status),
// Background mode is no longer supported, so there's nothing to cancel
error::bad_request(
"cancellation_not_supported",
"Background mode is not supported. Synchronous and streaming responses cannot be cancelled.",
)
}
}

2
sgl-model-gateway/src/routers/grpc/common/responses/mod.rs
View File

@@ -1,9 +1,11 @@
//! Shared response functionality used by both regular and harmony implementations
pub(crate) mod context;
pub(crate) mod handlers;
pub(crate) mod streaming;
pub(crate) mod utils;
// Re-export commonly used items
pub(crate) use context::ResponsesContext;
pub(crate) use streaming::build_sse_response;
pub(crate) use utils::{ensure_mcp_connection, persist_response_if_needed};

4
sgl-model-gateway/src/routers/grpc/harmony/mod.rs
View File

@@ -43,8 +43,6 @@ pub(crate) use builder::HarmonyBuilder;
pub(crate) use detector::HarmonyDetector;
pub(crate) use parser::HarmonyParserAdapter;
pub(crate) use processor::{HarmonyResponseProcessor, ResponsesIterationResult};
pub(crate) use responses::{
serve_harmony_responses, serve_harmony_responses_stream, HarmonyResponsesContext,
};
pub(crate) use responses::{serve_harmony_responses, serve_harmony_responses_stream};
pub(crate) use streaming::HarmonyStreamingProcessor;
pub(crate) use types::HarmonyMessage;

6
sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs
View File

@@ -5,7 +5,7 @@ use serde_json::{from_value, json, to_string, Value};
use tracing::{debug, error, warn};
use uuid::Uuid;
use super::{context::HarmonyResponsesContext, execution::ToolResult};
use super::execution::ToolResult;
use crate::{
data_connector::ResponseId,
mcp,
@@ -17,7 +17,7 @@ use crate::{
ResponsesRequest, ResponsesResponse, StringOrContentParts,
},
},
routers::error,
routers::{error, grpc::common::responses::ResponsesContext},
};
/// Record of a single MCP tool call execution
@@ -271,7 +271,7 @@ pub(super) fn inject_mcp_metadata(
/// If the request has `previous_response_id`, loads the response chain from storage
/// and prepends the conversation history to the request input items.
pub(super) async fn load_previous_messages(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
request: ResponsesRequest,
) -> Result<ResponsesRequest, Response> {
let Some(ref prev_id_str) = request.previous_response_id else {

3
sgl-model-gateway/src/routers/grpc/harmony/responses/mod.rs
View File

@@ -13,20 +13,17 @@
//!
//! ## Module Structure
//!
//! - `context` - HarmonyResponsesContext
//! - `non_streaming` - Non-streaming entry point and tool loop
//! - `streaming` - Streaming entry point and tool loop
//! - `execution` - MCP tool execution logic
//! - `common` - Shared helpers and state tracking
pub(crate) mod common;
pub(crate) mod context;
pub(crate) mod execution;
pub(crate) mod non_streaming;
pub(crate) mod streaming;
// Re-export types accessed via harmony::responses::TypeName
pub(crate) use context::HarmonyResponsesContext;
pub(crate) use execution::ToolResult;
pub(crate) use non_streaming::serve_harmony_responses;
pub(crate) use streaming::serve_harmony_responses_stream;

11
sgl-model-gateway/src/routers/grpc/harmony/responses/non_streaming.rs
View File

@@ -14,7 +14,6 @@ use super::{
build_mcp_tool_names_set, build_next_request_with_tools, inject_mcp_metadata,
load_previous_messages, McpCallTracking,
},
context::HarmonyResponsesContext,
execution::{convert_mcp_tools_to_response_tools, execute_mcp_tools, ToolResult},
};
use crate::{
@@ -29,7 +28,9 @@ use crate::{
routers::{
error,
grpc::{
common::responses::{ensure_mcp_connection, persist_response_if_needed},
common::responses::{
ensure_mcp_connection, persist_response_if_needed, ResponsesContext,
},
harmony::processor::ResponsesIterationResult,
},
mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS},
@@ -47,7 +48,7 @@ use crate::{
/// - Repeat from step 1 (full pipeline re-execution)
/// 4. If no tool calls, return final response
pub(crate) async fn serve_harmony_responses(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
request: ResponsesRequest,
) -> Result<ResponsesResponse, Response> {
// Clone request for persistence
@@ -90,7 +91,7 @@ pub(crate) async fn serve_harmony_responses(
///
/// Automatically executes MCP tools in a loop until no more tool calls or max iterations
async fn execute_with_mcp_loop(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
mut current_request: ResponsesRequest,
) -> Result<ResponsesResponse, Response> {
let mut iteration_count = 0;
@@ -319,7 +320,7 @@ async fn execute_with_mcp_loop(
///
/// For function tools or no tools - executes pipeline once and returns
async fn execute_without_mcp_loop(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
current_request: ResponsesRequest,
) -> Result<ResponsesResponse, Response> {
debug!("Executing Harmony Responses without MCP loop");

8
sgl-model-gateway/src/routers/grpc/harmony/responses/streaming.rs
View File

@@ -14,7 +14,6 @@ use super::{
build_mcp_tool_names_set, build_next_request_with_tools, load_previous_messages,
McpCallTracking,
},
context::HarmonyResponsesContext,
execution::{convert_mcp_tools_to_response_tools, execute_mcp_tools},
};
use crate::{
@@ -25,6 +24,7 @@ use crate::{
common::responses::{
build_sse_response, ensure_mcp_connection, persist_response_if_needed,
streaming::{OutputItemType, ResponseStreamEventEmitter},
ResponsesContext,
},
harmony::{processor::ResponsesIterationResult, streaming::HarmonyStreamingProcessor},
},
@@ -37,7 +37,7 @@ use crate::{
/// This is the streaming equivalent of `serve_harmony_responses()`.
/// Emits SSE events for lifecycle, MCP list_tools, and per-iteration streaming.
pub(crate) async fn serve_harmony_responses_stream(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
request: ResponsesRequest,
) -> Response {
// Load previous conversation history if previous_response_id is set
@@ -112,7 +112,7 @@ pub(crate) async fn serve_harmony_responses_stream(
/// - Emits final response.completed event
/// - Persists response internally
async fn execute_mcp_tool_loop_streaming(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
mut current_request: ResponsesRequest,
original_request: &ResponsesRequest,
emitter: &mut ResponseStreamEventEmitter,
@@ -459,7 +459,7 @@ async fn execute_mcp_tool_loop_streaming(
/// For function tools or no tools - executes pipeline once and emits completion.
/// The streaming processor handles all output items (reasoning, message, function tool calls).
async fn execute_without_mcp_streaming(
ctx: &HarmonyResponsesContext,
ctx: &ResponsesContext,
current_request: &ResponsesRequest,
original_request: &ResponsesRequest,
emitter: &mut ResponseStreamEventEmitter,

6
sgl-model-gateway/src/routers/grpc/pipeline.rs
View File

@@ -11,7 +11,7 @@ use tracing::{debug, error};
// Import embedding-specific and classify-specific stages
use super::regular::stages::classify::ClassifyResponseProcessingStage;
use super::{
common::stages::*,
common::{responses::ResponsesContext, stages::*},
context::*,
harmony,
regular::{
@@ -753,7 +753,7 @@ impl RequestPipeline {
pub async fn execute_harmony_responses(
&self,
request: &crate::protocols::responses::ResponsesRequest,
harmony_ctx: &harmony::responses::HarmonyResponsesContext,
harmony_ctx: &ResponsesContext,
) -> Result<harmony::ResponsesIterationResult, Response> {
// Create RequestContext for this Responses request
let mut ctx = RequestContext::for_responses(
@@ -816,7 +816,7 @@ impl RequestPipeline {
pub async fn execute_harmony_responses_streaming(
&self,
request: &crate::protocols::responses::ResponsesRequest,
harmony_ctx: &harmony::responses::HarmonyResponsesContext,
harmony_ctx: &ResponsesContext,
) -> Result<(ExecutionResult, Option<LoadGuards>), Response> {
// Create RequestContext for this Responses request
let mut ctx = RequestContext::for_responses(

3
sgl-model-gateway/src/routers/grpc/regular/responses/common.rs
View File

@@ -13,7 +13,6 @@ use serde_json::{json, Value};
use tracing::{debug, warn};
use uuid::Uuid;
use super::context::ResponsesContext;
use crate::{
data_connector::{self, ConversationId, ResponseId},
mcp::{self, McpManager},
@@ -25,7 +24,7 @@ use crate::{
ResponseOutputItem, ResponsesRequest,
},
},
routers::error,
routers::{error, grpc::common::responses::ResponsesContext},
};
// ============================================================================

84
sgl-model-gateway/src/routers/grpc/regular/responses/context.rs
View File

@@ -1,84 +0,0 @@
//! Context and types for /v1/responses endpoint handlers
//!
//! Bundles all dependencies needed by responses handlers to avoid passing
//! 10+ parameters to every function.
use std::{
collections::HashMap,
sync::{Arc, RwLock as StdRwLock},
};
use tokio::{sync::RwLock, task::JoinHandle};
use crate::{
data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage},
grpc_client::SglangSchedulerClient,
mcp::McpManager,
routers::grpc::{context::SharedComponents, pipeline::RequestPipeline},
};
/// Information stored for background tasks to enable end-to-end cancellation
///
/// This struct enables cancelling both the Rust task AND the Python scheduler processing.
/// The client field is lazily initialized during pipeline execution.
pub(crate) struct BackgroundTaskInfo {
/// Tokio task handle for aborting the Rust task
pub handle: JoinHandle<()>,
/// gRPC request_id sent to Python scheduler (chatcmpl-* prefix)
pub grpc_request_id: String,
/// gRPC client for sending abort requests to Python (set after client acquisition)
pub client: Arc<RwLock<Option<SglangSchedulerClient>>>,
}
/// Context for /v1/responses endpoint
///
/// All fields are Arc/shared references, so cloning this context is cheap.
#[derive(Clone)]
pub(crate) struct ResponsesContext {
/// Chat pipeline for executing requests
pub pipeline: Arc<RequestPipeline>,
/// Shared components (tokenizer, parsers)
pub components: Arc<SharedComponents>,
/// Response storage backend
pub response_storage: Arc<dyn ResponseStorage>,
/// Conversation storage backend
pub conversation_storage: Arc<dyn ConversationStorage>,
/// Conversation item storage backend
pub conversation_item_storage: Arc<dyn ConversationItemStorage>,
/// MCP manager for tool support
pub mcp_manager: Arc<McpManager>,
/// Server keys for MCP tools requested in this context
pub requested_servers: Arc<StdRwLock<Vec<String>>>,
/// Background task handles for cancellation support
pub background_tasks: Arc<RwLock<HashMap<String, BackgroundTaskInfo>>>,
}
impl ResponsesContext {
/// Create a new responses context
pub fn new(
pipeline: Arc<RequestPipeline>,
components: Arc<SharedComponents>,
response_storage: Arc<dyn ResponseStorage>,
conversation_storage: Arc<dyn ConversationStorage>,
conversation_item_storage: Arc<dyn ConversationItemStorage>,
mcp_manager: Arc<McpManager>,
) -> Self {
Self {
pipeline,
components,
response_storage,
conversation_storage,
conversation_item_storage,
mcp_manager,
requested_servers: Arc::new(StdRwLock::new(Vec::new())),
background_tasks: Arc::new(RwLock::new(HashMap::new())),
}
}
}

10
sgl-model-gateway/src/routers/grpc/regular/responses/handlers.rs
View File

@@ -34,13 +34,13 @@ use axum::{
use tracing::debug;
use uuid::Uuid;
use super::{
common::load_conversation_history, context::ResponsesContext, conversions, non_streaming,
streaming,
};
use super::{common::load_conversation_history, conversions, non_streaming, streaming};
use crate::{
protocols::responses::ResponsesRequest,
routers::{error, grpc::common::responses::ensure_mcp_connection},
routers::{
error,
grpc::common::responses::{ensure_mcp_connection, ResponsesContext},
},
};
/// Main handler for POST /v1/responses

3
sgl-model-gateway/src/routers/grpc/regular/responses/mod.rs
View File

@@ -9,15 +9,12 @@
//! - `streaming` - Streaming execution with MCP tool loop
//! - `common` - Shared helpers: ToolLoopState, tool preparation, MCP metadata builders
//! - `conversions` - Request/response conversion between Responses and Chat formats
//! - `context` - ResponsesContext and BackgroundTaskInfo
mod common;
mod context;
mod conversions;
mod handlers;
mod non_streaming;
mod streaming;
// Public exports
pub(crate) use context::ResponsesContext;
pub(crate) use handlers::route_responses;

5
sgl-model-gateway/src/routers/grpc/regular/responses/non_streaming.rs
View File

@@ -17,7 +17,6 @@ use super::{
extract_all_tool_calls_from_chat, load_conversation_history, prepare_chat_tools_and_choice,
ToolLoopState,
},
context::ResponsesContext,
conversions,
};
use crate::{
@@ -25,7 +24,9 @@ use crate::{
protocols::responses::{ResponseStatus, ResponsesRequest, ResponsesResponse},
routers::{
error,
grpc::common::responses::{ensure_mcp_connection, persist_response_if_needed},
grpc::common::responses::{
ensure_mcp_connection, persist_response_if_needed, ResponsesContext,
},
mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS},
},
};

2
sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs
View File

@@ -29,7 +29,6 @@ use super::{
build_next_request, convert_mcp_tools_to_chat_tools, extract_all_tool_calls_from_chat,
prepare_chat_tools_and_choice, ToolLoopState,
},
context::ResponsesContext,
conversions,
};
use crate::{
@@ -50,6 +49,7 @@ use crate::{
grpc::common::responses::{
build_sse_response, persist_response_if_needed,
streaming::{OutputItemType, ResponseStreamEventEmitter},
ResponsesContext,
},
mcp_utils::{extract_server_label, DEFAULT_MAX_ITERATIONS},
},

16
sgl-model-gateway/src/routers/grpc/router.rs
View File

@@ -11,12 +11,10 @@ use super::{
common::responses::{
handlers::{cancel_response_impl, get_response_impl},
utils::validate_worker_availability,
ResponsesContext,
},
context::SharedComponents,
harmony::{
serve_harmony_responses, serve_harmony_responses_stream, HarmonyDetector,
HarmonyResponsesContext,
},
harmony::{serve_harmony_responses, serve_harmony_responses_stream, HarmonyDetector},
pipeline::RequestPipeline,
regular::responses,
};
@@ -44,8 +42,8 @@ pub struct GrpcRouter {
embedding_pipeline: RequestPipeline,
classify_pipeline: RequestPipeline,
shared_components: Arc<SharedComponents>,
responses_context: responses::ResponsesContext,
harmony_responses_context: responses::ResponsesContext,
responses_context: ResponsesContext,
harmony_responses_context: ResponsesContext,
retry_config: RetryConfig,
}
@@ -113,7 +111,7 @@ impl GrpcRouter {
// Helper closure to create responses context with a given pipeline
let create_responses_context = |pipeline: &RequestPipeline| {
responses::ResponsesContext::new(
ResponsesContext::new(
Arc::new(pipeline.clone()),
shared_components.clone(),
ctx.response_storage.clone(),
@@ -286,15 +284,15 @@ impl GrpcRouter {
model_id.unwrap_or(UNKNOWN_MODEL_ID),
body.stream.unwrap_or(false)
);
let harmony_ctx = HarmonyResponsesContext::new(
let harmony_ctx = ResponsesContext::new(
Arc::new(self.harmony_pipeline.clone()),
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(),
self.harmony_responses_context.mcp_manager.clone(),
);
if body.stream.unwrap_or(false) {