From 51541404f84b3f344f1c80dcee0afaf82e0c10e6 Mon Sep 17 00:00:00 2001 From: Chang Su Date: Mon, 5 Jan 2026 11:39:33 -0800 Subject: [PATCH] [grpc] Refactor openai module (#16511) --- sgl-model-gateway/src/routers/openai/mod.rs | 10 +- .../src/routers/openai/provider.rs | 3 + .../openai/{ => responses}/accumulator.rs | 2 +- .../src/routers/openai/responses/common.rs | 113 ++++++++++++ .../src/routers/openai/{ => responses}/mcp.rs | 4 +- .../src/routers/openai/responses/mod.rs | 19 ++ .../routers/openai/responses/non_streaming.rs | 164 ++++++++++++++++++ .../openai/{ => responses}/streaming.rs | 117 +------------ .../openai/{ => responses}/tool_handler.rs | 6 +- .../{responses.rs => responses/utils.rs} | 2 + .../src/routers/openai/router.rs | 134 +------------- 11 files changed, 318 insertions(+), 256 deletions(-) rename sgl-model-gateway/src/routers/openai/{ => responses}/accumulator.rs (98%) create mode 100644 sgl-model-gateway/src/routers/openai/responses/common.rs rename sgl-model-gateway/src/routers/openai/{ => responses}/mcp.rs (99%) create mode 100644 sgl-model-gateway/src/routers/openai/responses/mod.rs create mode 100644 sgl-model-gateway/src/routers/openai/responses/non_streaming.rs rename sgl-model-gateway/src/routers/openai/{ => responses}/streaming.rs (90%) rename sgl-model-gateway/src/routers/openai/{ => responses}/tool_handler.rs (98%) rename sgl-model-gateway/src/routers/openai/{responses.rs => responses/utils.rs} (98%) diff --git a/sgl-model-gateway/src/routers/openai/mod.rs b/sgl-model-gateway/src/routers/openai/mod.rs index 02f270751..35a33fc08 100644 --- a/sgl-model-gateway/src/routers/openai/mod.rs +++ b/sgl-model-gateway/src/routers/openai/mod.rs @@ -7,15 +7,9 @@ //! - Multi-turn tool execution loops //! - SSE (Server-Sent Events) streaming -mod accumulator; mod context; -pub mod mcp; -pub mod provider; -mod responses; +mod provider; +pub mod responses; mod router; -mod streaming; -mod tool_handler; -// Re-export the main types for external use -pub use provider::{Provider, ProviderError, ProviderRegistry}; pub use router::OpenAIRouter; diff --git a/sgl-model-gateway/src/routers/openai/provider.rs b/sgl-model-gateway/src/routers/openai/provider.rs index 545555579..7539ec58d 100644 --- a/sgl-model-gateway/src/routers/openai/provider.rs +++ b/sgl-model-gateway/src/routers/openai/provider.rs @@ -218,6 +218,7 @@ impl ProviderRegistry { } } + #[allow(dead_code)] pub fn get(&self, provider_type: &ProviderType) -> &dyn Provider { self.providers .get(provider_type) @@ -232,6 +233,7 @@ impl ProviderRegistry { .unwrap_or_else(|| Arc::clone(&self.default_provider)) } + #[allow(dead_code)] pub fn get_for_model(&self, model_name: &str) -> &dyn Provider { match ProviderType::from_model_name(model_name) { Some(pt) => self.get(&pt), @@ -239,6 +241,7 @@ impl ProviderRegistry { } } + #[allow(dead_code)] pub fn default_provider(&self) -> &dyn Provider { self.default_provider.as_ref() } diff --git a/sgl-model-gateway/src/routers/openai/accumulator.rs b/sgl-model-gateway/src/routers/openai/responses/accumulator.rs similarity index 98% rename from sgl-model-gateway/src/routers/openai/accumulator.rs rename to sgl-model-gateway/src/routers/openai/responses/accumulator.rs index b9ed755aa..63f93c5b5 100644 --- a/sgl-model-gateway/src/routers/openai/accumulator.rs +++ b/sgl-model-gateway/src/routers/openai/responses/accumulator.rs @@ -3,7 +3,7 @@ use serde_json::Value; use tracing::warn; -use super::streaming::{extract_output_index, get_event_type}; +use super::common::{extract_output_index, get_event_type}; use crate::protocols::event_types::{OutputItemEvent, ResponseEvent}; // ============================================================================ diff --git a/sgl-model-gateway/src/routers/openai/responses/common.rs b/sgl-model-gateway/src/routers/openai/responses/common.rs new file mode 100644 index 000000000..b24e83bd3 --- /dev/null +++ b/sgl-model-gateway/src/routers/openai/responses/common.rs @@ -0,0 +1,113 @@ +//! Common SSE parsing and processing utilities for OpenAI responses +//! +//! This module contains shared helpers used by both streaming and accumulator modules. + +use std::borrow::Cow; + +use serde_json::Value; + +// ============================================================================ +// Helper Functions +// ============================================================================ + +/// Extract output_index from a JSON value +#[inline] +pub(super) fn extract_output_index(value: &Value) -> Option { + value.get("output_index")?.as_u64().map(|v| v as usize) +} + +/// Get event type from event name or parsed JSON, returning a reference to avoid allocation +#[inline] +pub(super) fn get_event_type<'a>(event_name: Option<&'a str>, parsed: &'a Value) -> &'a str { + event_name + .or_else(|| parsed.get("type").and_then(|v| v.as_str())) + .unwrap_or("") +} + +// ============================================================================ +// Chunk Processor +// ============================================================================ + +/// Processes incoming byte chunks into complete SSE blocks. +/// Handles buffering of partial chunks and CRLF normalization. +pub(super) struct ChunkProcessor { + pending: String, +} + +impl ChunkProcessor { + pub fn new() -> Self { + Self { + pending: String::new(), + } + } + + /// Append a chunk to the buffer, normalizing line endings + pub fn push_chunk(&mut self, chunk: &[u8]) { + let chunk_str = match std::str::from_utf8(chunk) { + Ok(s) => Cow::Borrowed(s), + Err(_) => Cow::Owned(String::from_utf8_lossy(chunk).into_owned()), + }; + // Normalize CRLF to LF without extra allocation + let mut chars = chunk_str.chars().peekable(); + while let Some(c) = chars.next() { + if c == '\r' && chars.peek() == Some(&'\n') { + // Skip \r when followed by \n + continue; + } + self.pending.push(c); + } + } + + /// Extract the next complete SSE block from the buffer, if available + pub fn next_block(&mut self) -> Option { + loop { + let pos = self.pending.find("\n\n")?; + let block = self.pending[..pos].to_string(); + self.pending.drain(..pos + 2); + + if !block.trim().is_empty() { + return Some(block); + } + // If block is empty, loop again to find the next one + } + } + + /// Check if there's remaining content in the buffer + pub fn has_remaining(&self) -> bool { + !self.pending.trim().is_empty() + } + + /// Take any remaining content from the buffer + pub fn take_remaining(&mut self) -> String { + std::mem::take(&mut self.pending) + } +} + +// ============================================================================ +// SSE Parsing +// ============================================================================ + +/// Parse an SSE block into event name and data +/// +/// Returns borrowed strings when possible to avoid allocations in hot paths. +/// Only allocates when multiple data lines need to be joined. +pub(super) fn parse_sse_block(block: &str) -> (Option<&str>, Cow<'_, str>) { + let mut event_name: Option<&str> = None; + let mut data_lines: Vec<&str> = Vec::new(); + + for line in block.lines() { + if let Some(rest) = line.strip_prefix("event:") { + event_name = Some(rest.trim()); + } else if let Some(rest) = line.strip_prefix("data:") { + data_lines.push(rest.trim_start()); + } + } + + let data = if data_lines.len() == 1 { + Cow::Borrowed(data_lines[0]) + } else { + Cow::Owned(data_lines.join("\n")) + }; + + (event_name, data) +} diff --git a/sgl-model-gateway/src/routers/openai/mcp.rs b/sgl-model-gateway/src/routers/openai/responses/mcp.rs similarity index 99% rename from sgl-model-gateway/src/routers/openai/mcp.rs rename to sgl-model-gateway/src/routers/openai/responses/mcp.rs index 95d07597f..4281f9934 100644 --- a/sgl-model-gateway/src/routers/openai/mcp.rs +++ b/sgl-model-gateway/src/routers/openai/responses/mcp.rs @@ -33,7 +33,7 @@ use crate::{ // ============================================================================ /// State for tracking multi-turn tool calling loop -pub(crate) struct ToolLoopState { +pub(super) struct ToolLoopState { /// Current iteration number (starts at 0, increments with each tool call) pub iteration: usize, /// Total number of tool calls executed @@ -83,7 +83,7 @@ impl ToolLoopState { /// Represents a function call being accumulated across delta events #[derive(Debug, Clone)] -pub(crate) struct FunctionCallInProgress { +pub(super) struct FunctionCallInProgress { pub call_id: String, pub name: String, pub arguments_buffer: String, diff --git a/sgl-model-gateway/src/routers/openai/responses/mod.rs b/sgl-model-gateway/src/routers/openai/responses/mod.rs new file mode 100644 index 000000000..5d61cf68d --- /dev/null +++ b/sgl-model-gateway/src/routers/openai/responses/mod.rs @@ -0,0 +1,19 @@ +//! OpenAI-compatible responses handling module +//! +//! This module provides comprehensive support for OpenAI Responses API with: +//! - Streaming and non-streaming response handling +//! - MCP (Model Context Protocol) tool interception and execution +//! - SSE (Server-Sent Events) parsing and forwarding +//! - Response accumulation for persistence +//! - Tool call detection and output index remapping + +mod accumulator; +mod common; +mod mcp; +mod non_streaming; +mod streaming; +mod tool_handler; +mod utils; + +pub use non_streaming::handle_non_streaming_response; +pub use streaming::handle_streaming_response; diff --git a/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs b/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs new file mode 100644 index 000000000..77e2e3c39 --- /dev/null +++ b/sgl-model-gateway/src/routers/openai/responses/non_streaming.rs @@ -0,0 +1,164 @@ +//! Non-streaming response handling for OpenAI-compatible responses +//! +//! This module handles non-streaming Responses API requests with MCP tool support. + +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, + Json, +}; +use serde_json::{json, Value}; +use tracing::warn; + +use super::{ + mcp::{execute_tool_loop, prepare_mcp_payload_for_streaming}, + utils::{mask_tools_as_mcp, patch_streaming_response_json}, +}; +use crate::routers::{ + header_utils::{apply_provider_headers, extract_auth_header}, + mcp_utils::{ensure_request_mcp_client, McpLoopConfig}, + openai::context::{PayloadState, RequestContext}, + persistence_utils::persist_conversation_items, +}; + +/// Handle a non-streaming responses request +pub async fn handle_non_streaming_response(mut ctx: RequestContext) -> Response { + let payload_state = match ctx.state.payload.take() { + Some(ps) => ps, + None => { + return (StatusCode::INTERNAL_SERVER_ERROR, "Payload not prepared").into_response(); + } + }; + + let PayloadState { + json: mut payload, + url, + previous_response_id, + } = payload_state; + + let original_body = ctx.responses_request(); + let worker = match ctx.worker() { + Some(w) => w.clone(), + None => { + return (StatusCode::INTERNAL_SERVER_ERROR, "Worker not selected").into_response(); + } + }; + let mcp_manager = match ctx.components.mcp_manager() { + Some(m) => m, + None => { + return (StatusCode::INTERNAL_SERVER_ERROR, "MCP manager required").into_response(); + } + }; + + if let Some(ref tools) = original_body.tools { + ensure_request_mcp_client(mcp_manager, tools.as_slice()).await; + } + + let active_mcp = if mcp_manager.list_tools().is_empty() { + None + } else { + Some(mcp_manager) + }; + + let mut response_json: Value; + + if let Some(mcp) = active_mcp { + let config = McpLoopConfig::default(); + prepare_mcp_payload_for_streaming(&mut payload, mcp); + + match execute_tool_loop( + ctx.components.client(), + &url, + ctx.headers(), + payload, + original_body, + mcp, + &config, + ) + .await + { + Ok(resp) => response_json = resp, + Err(err) => { + worker.circuit_breaker().record_failure(); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({"error": {"message": err}})), + ) + .into_response(); + } + } + } else { + let mut request_builder = ctx.components.client().post(&url).json(&payload); + let auth_header = extract_auth_header(ctx.headers(), worker.api_key()); + request_builder = apply_provider_headers(request_builder, &url, auth_header.as_ref()); + + let response = match request_builder.send().await { + Ok(r) => r, + Err(e) => { + worker.circuit_breaker().record_failure(); + tracing::error!( + url = %url, + error = %e, + "Failed to forward request to OpenAI" + ); + return ( + StatusCode::BAD_GATEWAY, + format!("Failed to forward request to OpenAI: {}", e), + ) + .into_response(); + } + }; + + if !response.status().is_success() { + worker.circuit_breaker().record_failure(); + let status = StatusCode::from_u16(response.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + let body = response.text().await.unwrap_or_default(); + return (status, body).into_response(); + } + + response_json = match response.json::().await { + Ok(r) => r, + Err(e) => { + worker.circuit_breaker().record_failure(); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to parse upstream response: {}", e), + ) + .into_response(); + } + }; + + worker.circuit_breaker().record_success(); + } + + mask_tools_as_mcp(&mut response_json, original_body); + patch_streaming_response_json( + &mut response_json, + original_body, + previous_response_id.as_deref(), + ); + + if let Err(err) = persist_conversation_items( + ctx.components + .conversation_storage() + .expect("Conversation storage required") + .clone(), + ctx.components + .conversation_item_storage() + .expect("Conversation item storage required") + .clone(), + ctx.components + .response_storage() + .expect("Response storage required") + .clone(), + &response_json, + original_body, + ) + .await + { + warn!("Failed to persist conversation items: {}", err); + } + + (StatusCode::OK, Json(response_json)).into_response() +} diff --git a/sgl-model-gateway/src/routers/openai/streaming.rs b/sgl-model-gateway/src/routers/openai/responses/streaming.rs similarity index 90% rename from sgl-model-gateway/src/routers/openai/streaming.rs rename to sgl-model-gateway/src/routers/openai/responses/streaming.rs index 380c9c265..2b012cbdf 100644 --- a/sgl-model-gateway/src/routers/openai/streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/streaming.rs @@ -21,16 +21,15 @@ use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::warn; -// Import from sibling modules -use super::accumulator::StreamingResponseAccumulator; use super::{ - context::{RequestContext, StreamingEventContext, StreamingRequest}, + accumulator::StreamingResponseAccumulator, + common::{extract_output_index, get_event_type, parse_sse_block, ChunkProcessor}, mcp::{ 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}, + utils::{mask_tools_as_mcp, patch_streaming_response_json, rewrite_streaming_block}, }; use crate::{ protocols::{ @@ -43,116 +42,11 @@ use crate::{ routers::{ header_utils::{apply_request_headers, preserve_response_headers}, mcp_utils::{ensure_request_mcp_client, McpLoopConfig}, + openai::context::{RequestContext, StreamingEventContext, StreamingRequest}, persistence_utils::persist_conversation_items, }, }; -// ============================================================================ -// Helper Functions -// ============================================================================ - -/// Extract output_index from a JSON value -#[inline] -pub(super) fn extract_output_index(value: &Value) -> Option { - value.get("output_index")?.as_u64().map(|v| v as usize) -} - -/// Get event type from event name or parsed JSON, returning a reference to avoid allocation -#[inline] -pub(super) fn get_event_type<'a>(event_name: Option<&'a str>, parsed: &'a Value) -> &'a str { - event_name - .or_else(|| parsed.get("type").and_then(|v| v.as_str())) - .unwrap_or("") -} - -// ============================================================================ -// Chunk Processor -// ============================================================================ - -/// Processes incoming byte chunks into complete SSE blocks. -/// Handles buffering of partial chunks and CRLF normalization. -pub(super) struct ChunkProcessor { - pending: String, -} - -impl ChunkProcessor { - pub fn new() -> Self { - Self { - pending: String::new(), - } - } - - /// Append a chunk to the buffer, normalizing line endings - pub fn push_chunk(&mut self, chunk: &[u8]) { - let chunk_str = match std::str::from_utf8(chunk) { - Ok(s) => Cow::Borrowed(s), - Err(_) => Cow::Owned(String::from_utf8_lossy(chunk).into_owned()), - }; - // Normalize CRLF to LF without extra allocation - let mut chars = chunk_str.chars().peekable(); - while let Some(c) = chars.next() { - if c == '\r' && chars.peek() == Some(&'\n') { - // Skip \r when followed by \n - continue; - } - self.pending.push(c); - } - } - - /// Extract the next complete SSE block from the buffer, if available - pub fn next_block(&mut self) -> Option { - loop { - let pos = self.pending.find("\n\n")?; - let block = self.pending[..pos].to_string(); - self.pending.drain(..pos + 2); - - if !block.trim().is_empty() { - return Some(block); - } - // If block is empty, loop again to find the next one - } - } - - /// Check if there's remaining content in the buffer - pub fn has_remaining(&self) -> bool { - !self.pending.trim().is_empty() - } - - /// Take any remaining content from the buffer - pub fn take_remaining(&mut self) -> String { - std::mem::take(&mut self.pending) - } -} - -// ============================================================================ -// SSE Parsing -// ============================================================================ - -/// Parse an SSE block into event name and data -/// -/// Returns borrowed strings when possible to avoid allocations in hot paths. -/// Only allocates when multiple data lines need to be joined. -pub(super) fn parse_sse_block(block: &str) -> (Option<&str>, Cow<'_, str>) { - let mut event_name: Option<&str> = None; - let mut data_lines: Vec<&str> = Vec::new(); - - for line in block.lines() { - if let Some(rest) = line.strip_prefix("event:") { - event_name = Some(rest.trim()); - } else if let Some(rest) = line.strip_prefix("data:") { - data_lines.push(rest.trim_start()); - } - } - - let data = if data_lines.len() == 1 { - Cow::Borrowed(data_lines[0]) - } else { - Cow::Owned(data_lines.join("\n")) - }; - - (event_name, data) -} - // ============================================================================ // Event Transformation and Forwarding // ============================================================================ @@ -1075,7 +969,8 @@ pub(super) async fn handle_streaming_with_tool_interception( response } -pub(super) async fn handle_streaming_response(ctx: RequestContext) -> Response { +/// Main entry point for streaming responses +pub async fn handle_streaming_response(ctx: RequestContext) -> Response { let worker = ctx.worker().expect("Worker not selected").clone(); let circuit_breaker = worker.circuit_breaker(); let headers = ctx.headers().cloned(); diff --git a/sgl-model-gateway/src/routers/openai/tool_handler.rs b/sgl-model-gateway/src/routers/openai/responses/tool_handler.rs similarity index 98% rename from sgl-model-gateway/src/routers/openai/tool_handler.rs rename to sgl-model-gateway/src/routers/openai/responses/tool_handler.rs index 5c9128461..7d358d54a 100644 --- a/sgl-model-gateway/src/routers/openai/tool_handler.rs +++ b/sgl-model-gateway/src/routers/openai/responses/tool_handler.rs @@ -7,8 +7,8 @@ use tracing::warn; use super::{ accumulator::StreamingResponseAccumulator, + common::{extract_output_index, get_event_type}, mcp::FunctionCallInProgress, - streaming::{extract_output_index, get_event_type}, }; use crate::protocols::event_types::{ is_function_call_type, FunctionCallEvent, OutputItemEvent, ResponseEvent, @@ -20,7 +20,7 @@ use crate::protocols::event_types::{ /// Action to take based on streaming event processing #[derive(Debug)] -pub(crate) enum StreamAction { +pub(super) enum StreamAction { Forward, // Pass event to client Buffer, // Accumulate for tool execution ExecuteTools, // Function call complete, execute now @@ -32,7 +32,7 @@ pub(crate) enum StreamAction { /// Maps upstream output indices to sequential downstream indices #[derive(Debug, Default)] -pub(crate) struct OutputIndexMapper { +pub(super) struct OutputIndexMapper { next_index: usize, // Map upstream output_index -> remapped output_index assigned: HashMap, diff --git a/sgl-model-gateway/src/routers/openai/responses.rs b/sgl-model-gateway/src/routers/openai/responses/utils.rs similarity index 98% rename from sgl-model-gateway/src/routers/openai/responses.rs rename to sgl-model-gateway/src/routers/openai/responses/utils.rs index 8c6ee1fac..f3b51ef5b 100644 --- a/sgl-model-gateway/src/routers/openai/responses.rs +++ b/sgl-model-gateway/src/routers/openai/responses/utils.rs @@ -1,3 +1,5 @@ +//! Response patching and transformation utilities for OpenAI responses + use serde_json::{json, Map, Value}; use tracing::warn; diff --git a/sgl-model-gateway/src/routers/openai/router.rs b/sgl-model-gateway/src/routers/openai/router.rs index e62765a51..336679318 100644 --- a/sgl-model-gateway/src/routers/openai/router.rs +++ b/sgl-model-gateway/src/routers/openai/router.rs @@ -23,10 +23,8 @@ use super::{ ComponentRefs, PayloadState, RequestContext, ResponsesComponents, SharedComponents, WorkerSelection, }, - mcp::{execute_tool_loop, prepare_mcp_payload_for_streaming}, provider::ProviderRegistry, - responses::{mask_tools_as_mcp, patch_streaming_response_json}, - streaming::handle_streaming_response, + responses::{handle_non_streaming_response, handle_streaming_response}, }; use crate::{ app_context::AppContext, @@ -44,11 +42,7 @@ use crate::{ ResponsesGetParams, ResponsesRequest, }, }, - routers::{ - header_utils::{apply_provider_headers, extract_auth_header}, - mcp_utils::{ensure_request_mcp_client, McpLoopConfig}, - persistence_utils::persist_conversation_items, - }, + routers::header_utils::{apply_provider_headers, extract_auth_header}, }; pub struct OpenAIRouter { @@ -367,128 +361,6 @@ impl OpenAIRouter { } } } - - async fn handle_non_streaming_response(&self, mut ctx: RequestContext) -> Response { - let payload_state = ctx.take_payload().expect("Payload not prepared"); - let mut payload = payload_state.json; - let url = payload_state.url; - let previous_response_id = payload_state.previous_response_id; - let original_body = ctx.responses_request(); - let worker = ctx.worker().expect("Worker not selected"); - let mcp_manager = ctx.components.mcp_manager().expect("MCP manager required"); - - if let Some(ref tools) = original_body.tools { - ensure_request_mcp_client(mcp_manager, tools.as_slice()).await; - } - - let active_mcp = if mcp_manager.list_tools().is_empty() { - None - } else { - Some(mcp_manager) - }; - - let mut response_json: Value; - - if let Some(mcp) = active_mcp { - let config = McpLoopConfig::default(); - prepare_mcp_payload_for_streaming(&mut payload, mcp); - - match execute_tool_loop( - ctx.components.client(), - &url, - ctx.headers(), - payload, - original_body, - mcp, - &config, - ) - .await - { - Ok(resp) => response_json = resp, - Err(err) => { - worker.circuit_breaker().record_failure(); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(json!({"error": {"message": err}})), - ) - .into_response(); - } - } - } else { - let mut request_builder = ctx.components.client().post(&url).json(&payload); - let auth_header = extract_auth_header(ctx.headers(), worker.api_key()); - request_builder = apply_provider_headers(request_builder, &url, auth_header.as_ref()); - - let response = match request_builder.send().await { - Ok(r) => r, - Err(e) => { - worker.circuit_breaker().record_failure(); - tracing::error!( - url = %url, - error = %e, - "Failed to forward request to OpenAI" - ); - return ( - StatusCode::BAD_GATEWAY, - format!("Failed to forward request to OpenAI: {}", e), - ) - .into_response(); - } - }; - - if !response.status().is_success() { - worker.circuit_breaker().record_failure(); - let status = StatusCode::from_u16(response.status().as_u16()) - .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - let body = response.text().await.unwrap_or_default(); - return (status, body).into_response(); - } - - response_json = match response.json::().await { - Ok(r) => r, - Err(e) => { - worker.circuit_breaker().record_failure(); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Failed to parse upstream response: {}", e), - ) - .into_response(); - } - }; - - worker.circuit_breaker().record_success(); - } - - mask_tools_as_mcp(&mut response_json, original_body); - patch_streaming_response_json( - &mut response_json, - original_body, - previous_response_id.as_deref(), - ); - - if let Err(err) = persist_conversation_items( - ctx.components - .conversation_storage() - .expect("Conversation storage required") - .clone(), - ctx.components - .conversation_item_storage() - .expect("Conversation item storage required") - .clone(), - ctx.components - .response_storage() - .expect("Response storage required") - .clone(), - &response_json, - original_body, - ) - .await - { - warn!("Failed to persist conversation items: {}", err); - } - - (StatusCode::OK, Json(response_json)).into_response() - } } #[async_trait::async_trait] @@ -1078,7 +950,7 @@ impl crate::routers::RouterTrait for OpenAIRouter { let response = if ctx.is_streaming() { handle_streaming_response(ctx).await } else { - self.handle_non_streaming_response(ctx).await + handle_non_streaming_response(ctx).await }; // Record duration only for successful requests (errors tracked inside handlers)