[grpc] Refactor openai module (#16511)

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;

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()
     }

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};
 
 // ============================================================================

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<usize> {
+    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<String> {
+        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)
+}

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,

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;

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::<Value>().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()
+}

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<usize> {
-    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<String> {
-        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();

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<usize, usize>,

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;
 

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::<Value>().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)