diff --git a/crates/mcp/src/core/config.rs b/crates/mcp/src/core/config.rs index c50f97a8cc..9de3900983 100644 --- a/crates/mcp/src/core/config.rs +++ b/crates/mcp/src/core/config.rs @@ -302,6 +302,8 @@ pub enum BuiltinToolType { CodeInterpreter, /// File search tool (OpenAI: file_search) FileSearch, + /// Image generation tool (OpenAI: image_generation) + ImageGeneration, } impl BuiltinToolType { @@ -311,6 +313,7 @@ impl BuiltinToolType { BuiltinToolType::WebSearchPreview => ResponseFormatConfig::WebSearchCall, BuiltinToolType::CodeInterpreter => ResponseFormatConfig::CodeInterpreterCall, BuiltinToolType::FileSearch => ResponseFormatConfig::FileSearchCall, + BuiltinToolType::ImageGeneration => ResponseFormatConfig::ImageGenerationCall, } } } @@ -321,6 +324,7 @@ impl fmt::Display for BuiltinToolType { BuiltinToolType::WebSearchPreview => write!(f, "web_search_preview"), BuiltinToolType::CodeInterpreter => write!(f, "code_interpreter"), BuiltinToolType::FileSearch => write!(f, "file_search"), + BuiltinToolType::ImageGeneration => write!(f, "image_generation"), } } } @@ -350,6 +354,7 @@ pub enum ResponseFormatConfig { WebSearchCall, CodeInterpreterCall, FileSearchCall, + ImageGenerationCall, } /// Argument mapping configuration for tool aliases. @@ -1025,6 +1030,10 @@ tools: "\"code_interpreter_call\"", ), (ResponseFormatConfig::FileSearchCall, "\"file_search_call\""), + ( + ResponseFormatConfig::ImageGenerationCall, + "\"image_generation_call\"", + ), ]; for (format, expected) in formats { @@ -1191,6 +1200,7 @@ policy: (BuiltinToolType::WebSearchPreview, "\"web_search_preview\""), (BuiltinToolType::CodeInterpreter, "\"code_interpreter\""), (BuiltinToolType::FileSearch, "\"file_search\""), + (BuiltinToolType::ImageGeneration, "\"image_generation\""), ]; for (builtin_type, expected) in types { @@ -1216,6 +1226,10 @@ policy: BuiltinToolType::FileSearch.response_format(), ResponseFormatConfig::FileSearchCall ); + assert_eq!( + BuiltinToolType::ImageGeneration.response_format(), + ResponseFormatConfig::ImageGenerationCall + ); } #[test] diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index b058ee5633..3f6973bc84 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -57,5 +57,6 @@ pub use responses_bridge::{ pub use tenant::{SessionId, TenantContext, TenantId}; // Re-export from transform pub use transform::{ - extract_embedded_openai_responses, mcp_response_item_id, ResponseFormat, ResponseTransformer, + compact_image_generation_output, extract_embedded_openai_responses, mcp_response_item_id, + ResponseFormat, ResponseTransformer, }; diff --git a/crates/mcp/src/transform/mod.rs b/crates/mcp/src/transform/mod.rs index a3921af07f..6127c7383d 100644 --- a/crates/mcp/src/transform/mod.rs +++ b/crates/mcp/src/transform/mod.rs @@ -23,6 +23,7 @@ mod transformer; mod types; pub use transformer::{ - extract_embedded_openai_responses, mcp_response_item_id, ResponseTransformer, + compact_image_generation_output, extract_embedded_openai_responses, mcp_response_item_id, + ResponseTransformer, }; pub use types::ResponseFormat; diff --git a/crates/mcp/src/transform/transformer.rs b/crates/mcp/src/transform/transformer.rs index 0cc2121fb2..414884af22 100644 --- a/crates/mcp/src/transform/transformer.rs +++ b/crates/mcp/src/transform/transformer.rs @@ -2,7 +2,8 @@ use openai_protocol::responses::{ CodeInterpreterCallStatus, CodeInterpreterOutput, FileSearchCallStatus, FileSearchResult, - ResponseOutputItem, WebSearchAction, WebSearchCallStatus, WebSearchSource, + ImageGenerationCallStatus, ResponseOutputItem, WebSearchAction, WebSearchCallStatus, + WebSearchSource, }; use tracing::warn; @@ -79,6 +80,9 @@ impl ResponseTransformer { Self::to_code_interpreter_call(result, tool_call_id) } ResponseFormat::FileSearchCall => Self::to_file_search_call(result, tool_call_id), + ResponseFormat::ImageGenerationCall => { + Self::to_image_generation_call(result, tool_call_id) + } } } @@ -236,6 +240,80 @@ impl ResponseTransformer { } } + /// Transform MCP image generation results to OpenAI image_generation_call format. + /// + /// Extracts fields from the MCP `CallToolResult`: + /// - `result` / `image_base64` / `b64_json`: base64-encoded image bytes + /// - `revised_prompt`: optional rewritten prompt preserved for replay + /// + /// Also probes embedded text blocks carrying `{"openai_response": {...}}` + /// payloads for the same fields, mirroring the pattern used by + /// `to_web_search_call`. + fn to_image_generation_call( + result: &serde_json::Value, + tool_call_id: &str, + ) -> ResponseOutputItem { + let (image_b64, revised_prompt) = Self::extract_image_generation_fields(result); + + ResponseOutputItem::ImageGenerationCall { + id: format!("ig_{tool_call_id}"), + // `result` is non-optional in the spec; fall back to an empty + // string when the MCP server returns no image data so the shape + // still matches `ResponseOutputItem::ImageGenerationCall`. Router + // wiring in R6.2/R6.3/R6.4 is responsible for surfacing an error + // status when `image_b64` is missing. + result: image_b64.unwrap_or_default(), + revised_prompt, + status: ImageGenerationCallStatus::Completed, + } + } + + /// Extract `(image_base64, revised_prompt)` from an MCP result payload. + fn extract_image_generation_fields( + result: &serde_json::Value, + ) -> (Option, Option) { + let mut image_b64: Option = None; + let mut revised_prompt: Option = None; + + // Merge helper: fills any still-unset output slot from an object. + // Using mutable accumulation (instead of early-returning on first + // match) prevents losing fields that an MCP server distributes + // across multiple text blocks — e.g. `revised_prompt` in one and + // `image_base64` in another. First occurrence wins for each slot. + let mut update_fields = |obj: &serde_json::Map| { + if image_b64.is_none() { + image_b64 = obj + .get("result") + .and_then(|v| v.as_str()) + .or_else(|| obj.get("image_base64").and_then(|v| v.as_str())) + .or_else(|| obj.get("b64_json").and_then(|v| v.as_str())) + .map(String::from); + } + if revised_prompt.is_none() { + revised_prompt = obj + .get("revised_prompt") + .and_then(|v| v.as_str()) + .map(String::from); + } + }; + + // 1) Direct object fields: { result | image_base64 | b64_json, revised_prompt } + if let Some(obj) = result.as_object() { + update_fields(obj); + } + + // 2) Embedded openai_response payloads inside MCP text blocks. + if result.is_array() { + for openai_response in extract_embedded_openai_responses(result) { + if let Some(obj) = openai_response.as_object() { + update_fields(obj); + } + } + } + + (image_b64, revised_prompt) + } + /// Extract web sources from MCP result. fn extract_web_sources(result: &serde_json::Value) -> Vec { let mut sources = Vec::new(); @@ -426,6 +504,27 @@ impl ResponseTransformer { } } +/// Strip the base64 `result` payload from an `ImageGenerationCall` output +/// item so stored multi-turn context does not balloon with large image bytes. +/// +/// Per OpenAI spec the `result` field is required on the wire when the item +/// is freshly emitted, but multi-turn replay/storage references the image by +/// `id` only (see `ResponseInputOutputItem::ImageGenerationCall` where +/// `result` is `Option`). This helper clears the payload for storage +/// while preserving `id`, `revised_prompt`, and `status`. +/// +/// For any non-image item this is a no-op. +pub fn compact_image_generation_output(item: &mut ResponseOutputItem) { + if let ResponseOutputItem::ImageGenerationCall { result, .. } = item { + // `result.clear()` zeros the length but keeps the heap buffer + // allocated — for base64 image bytes that can be several MB of + // wasted capacity per stored item, defeating the compaction goal. + // Replace with a fresh empty string to actually free the backing + // buffer. + *result = String::new(); + } +} + fn parse_text_block_payload(item: &serde_json::Value) -> Option { let Some(obj) = item.as_object() else { warn!("Expected MCP result item to be an object"); @@ -876,4 +975,230 @@ mod tests { _ => panic!("Expected FileSearchCall"), } } + + #[test] + fn test_image_generation_transform_direct_fields() { + let result = json!({ + "result": "BASE64_IMAGE_BYTES", + "revised_prompt": "a serene mountain at sunrise" + }); + + let transformed = ResponseTransformer::transform( + &result, + &ResponseFormat::ImageGenerationCall, + "req-img-1", + "server", + "image_generation", + r#"{"prompt":"mountain"}"#, + ); + + match transformed { + ResponseOutputItem::ImageGenerationCall { + id, + result, + revised_prompt, + status, + } => { + assert_eq!(id, "ig_req-img-1"); + assert_eq!(result, "BASE64_IMAGE_BYTES"); + assert_eq!( + revised_prompt, + Some("a serene mountain at sunrise".to_string()) + ); + assert_eq!(status, ImageGenerationCallStatus::Completed); + } + _ => panic!("Expected ImageGenerationCall"), + } + } + + #[test] + fn test_image_generation_transform_b64_json_alias() { + let result = json!({ + "b64_json": "PNG_BASE64", + }); + + let transformed = ResponseTransformer::transform( + &result, + &ResponseFormat::ImageGenerationCall, + "req-img-2", + "server", + "image_generation", + "{}", + ); + + match transformed { + ResponseOutputItem::ImageGenerationCall { + result, + revised_prompt, + .. + } => { + assert_eq!(result, "PNG_BASE64"); + assert!(revised_prompt.is_none()); + } + _ => panic!("Expected ImageGenerationCall"), + } + } + + #[test] + fn test_image_generation_transform_embedded_openai_response() { + let result = json!([ + { + "type": "text", + "text": r#"{ + "openai_response": { + "result": "EMBEDDED_BYTES", + "revised_prompt": "tweaked prompt" + } + }"# + } + ]); + + let transformed = ResponseTransformer::transform( + &result, + &ResponseFormat::ImageGenerationCall, + "req-img-3", + "server", + "image_generation", + "{}", + ); + + match transformed { + ResponseOutputItem::ImageGenerationCall { + result, + revised_prompt, + .. + } => { + assert_eq!(result, "EMBEDDED_BYTES"); + assert_eq!(revised_prompt, Some("tweaked prompt".to_string())); + } + _ => panic!("Expected ImageGenerationCall"), + } + } + + #[test] + fn test_image_generation_transform_fields_distributed_across_blocks() { + // MCP servers may split revised_prompt and image bytes across + // different text blocks. The extractor must accumulate fields + // from all blocks rather than returning early on the first hit. + let result = json!([ + { + "type": "text", + "text": r#"{ + "openai_response": { + "revised_prompt": "distributed prompt" + } + }"# + }, + { + "type": "text", + "text": r#"{ + "openai_response": { + "result": "DISTRIBUTED_BYTES" + } + }"# + } + ]); + + let transformed = ResponseTransformer::transform( + &result, + &ResponseFormat::ImageGenerationCall, + "req-img-dist", + "server", + "image_generation", + "{}", + ); + + match transformed { + ResponseOutputItem::ImageGenerationCall { + result, + revised_prompt, + .. + } => { + assert_eq!(result, "DISTRIBUTED_BYTES"); + assert_eq!(revised_prompt, Some("distributed prompt".to_string())); + } + _ => panic!("Expected ImageGenerationCall"), + } + } + + #[test] + fn test_image_generation_transform_missing_image_data() { + // When the MCP server returns neither image nor prompt, the transformer + // still produces a well-formed ImageGenerationCall with an empty result. + // Per-router wiring in R6.2/R6.3/R6.4 owns surfacing an error status. + let result = json!({}); + + let transformed = ResponseTransformer::transform( + &result, + &ResponseFormat::ImageGenerationCall, + "req-img-4", + "server", + "image_generation", + "{}", + ); + + match transformed { + ResponseOutputItem::ImageGenerationCall { + id, + result, + revised_prompt, + status, + } => { + assert_eq!(id, "ig_req-img-4"); + assert!(result.is_empty()); + assert!(revised_prompt.is_none()); + assert_eq!(status, ImageGenerationCallStatus::Completed); + } + _ => panic!("Expected ImageGenerationCall"), + } + } + + #[test] + fn test_compact_image_generation_output_strips_base64() { + let mut item = ResponseOutputItem::ImageGenerationCall { + id: "ig_abc".to_string(), + result: "A_VERY_LONG_BASE64_STRING".to_string(), + revised_prompt: Some("mountain".to_string()), + status: ImageGenerationCallStatus::Completed, + }; + + compact_image_generation_output(&mut item); + + match item { + ResponseOutputItem::ImageGenerationCall { + id, + result, + revised_prompt, + status, + } => { + assert_eq!(id, "ig_abc"); + assert!(result.is_empty()); + assert_eq!(revised_prompt, Some("mountain".to_string())); + assert_eq!(status, ImageGenerationCallStatus::Completed); + } + _ => panic!("Expected ImageGenerationCall"), + } + } + + #[test] + fn test_compact_image_generation_output_is_noop_for_other_items() { + let mut item = ResponseOutputItem::WebSearchCall { + id: "ws_xyz".to_string(), + status: WebSearchCallStatus::Completed, + action: WebSearchAction::Search { + query: Some("rust".to_string()), + queries: vec!["rust".to_string()], + sources: vec![], + }, + results: None, + }; + + compact_image_generation_output(&mut item); + + // WebSearchCall must be unchanged — no panic, id still present. + match item { + ResponseOutputItem::WebSearchCall { id, .. } => assert_eq!(id, "ws_xyz"), + _ => panic!("Expected WebSearchCall"), + } + } } diff --git a/crates/mcp/src/transform/types.rs b/crates/mcp/src/transform/types.rs index f508dc9d4b..7c489812ad 100644 --- a/crates/mcp/src/transform/types.rs +++ b/crates/mcp/src/transform/types.rs @@ -17,6 +17,8 @@ pub enum ResponseFormat { CodeInterpreterCall, /// Transform to OpenAI file_search_call format FileSearchCall, + /// Transform to OpenAI image_generation_call format + ImageGenerationCall, } impl From for ResponseFormat { @@ -26,6 +28,7 @@ impl From for ResponseFormat { ResponseFormatConfig::WebSearchCall => ResponseFormat::WebSearchCall, ResponseFormatConfig::CodeInterpreterCall => ResponseFormat::CodeInterpreterCall, ResponseFormatConfig::FileSearchCall => ResponseFormat::FileSearchCall, + ResponseFormatConfig::ImageGenerationCall => ResponseFormat::ImageGenerationCall, } } } @@ -44,6 +47,10 @@ mod tests { "\"code_interpreter_call\"", ), (ResponseFormat::FileSearchCall, "\"file_search_call\""), + ( + ResponseFormat::ImageGenerationCall, + "\"image_generation_call\"", + ), ]; for (format, expected) in formats { diff --git a/crates/protocols/src/event_types.rs b/crates/protocols/src/event_types.rs index fbf23e8c84..26ec7bc60f 100644 --- a/crates/protocols/src/event_types.rs +++ b/crates/protocols/src/event_types.rs @@ -263,6 +263,42 @@ impl fmt::Display for FileSearchCallEvent { } } +/// Image generation call events for streaming. +/// +/// Mirrors OpenAI Python SDK 2.8.1 `response_image_gen_call_*_event.py` and +/// `.claude/_audit/openai-responses-api-spec.md` §tools (image_generation). +/// `PartialImage` is emitted 0-3 times per call when the tool is configured +/// with `partial_images`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ImageGenerationCallEvent { + InProgress, + Generating, + PartialImage, + Completed, +} + +impl ImageGenerationCallEvent { + pub const IN_PROGRESS: &'static str = "response.image_generation_call.in_progress"; + pub const GENERATING: &'static str = "response.image_generation_call.generating"; + pub const PARTIAL_IMAGE: &'static str = "response.image_generation_call.partial_image"; + pub const COMPLETED: &'static str = "response.image_generation_call.completed"; + + pub const fn as_str(self) -> &'static str { + match self { + Self::InProgress => Self::IN_PROGRESS, + Self::Generating => Self::GENERATING, + Self::PartialImage => Self::PARTIAL_IMAGE, + Self::Completed => Self::COMPLETED, + } + } +} + +impl fmt::Display for ImageGenerationCallEvent { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + /// Item type discriminators used in output items #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum ItemType { @@ -274,6 +310,7 @@ pub enum ItemType { WebSearchCall, CodeInterpreterCall, FileSearchCall, + ImageGenerationCall, } impl ItemType { @@ -286,6 +323,7 @@ impl ItemType { pub const WEB_SEARCH_CALL: &'static str = "web_search_call"; pub const CODE_INTERPRETER_CALL: &'static str = "code_interpreter_call"; pub const FILE_SEARCH_CALL: &'static str = "file_search_call"; + pub const IMAGE_GENERATION_CALL: &'static str = "image_generation_call"; pub const fn as_str(self) -> &'static str { match self { @@ -297,6 +335,7 @@ impl ItemType { Self::WebSearchCall => Self::WEB_SEARCH_CALL, Self::CodeInterpreterCall => Self::CODE_INTERPRETER_CALL, Self::FileSearchCall => Self::FILE_SEARCH_CALL, + Self::ImageGenerationCall => Self::IMAGE_GENERATION_CALL, } } @@ -309,7 +348,10 @@ impl ItemType { pub const fn is_builtin_tool_call(self) -> bool { matches!( self, - Self::WebSearchCall | Self::CodeInterpreterCall | Self::FileSearchCall + Self::WebSearchCall + | Self::CodeInterpreterCall + | Self::FileSearchCall + | Self::ImageGenerationCall ) } } diff --git a/model_gateway/src/routers/common/mcp_utils.rs b/model_gateway/src/routers/common/mcp_utils.rs index 613c6abc81..a1f9067b72 100644 --- a/model_gateway/src/routers/common/mcp_utils.rs +++ b/model_gateway/src/routers/common/mcp_utils.rs @@ -158,6 +158,7 @@ pub fn collect_builtin_routing( let builtin_type = match tool { ResponseTool::WebSearchPreview(_) => BuiltinToolType::WebSearchPreview, ResponseTool::CodeInterpreter(_) => BuiltinToolType::CodeInterpreter, + ResponseTool::ImageGeneration(_) => BuiltinToolType::ImageGeneration, _ => continue, }; @@ -198,6 +199,7 @@ pub fn extract_builtin_types(tools: &[ResponseTool]) -> Vec { .filter_map(|t| match t { ResponseTool::WebSearchPreview(_) => Some(BuiltinToolType::WebSearchPreview), ResponseTool::CodeInterpreter(_) => Some(BuiltinToolType::CodeInterpreter), + ResponseTool::ImageGeneration(_) => Some(BuiltinToolType::ImageGeneration), _ => None, }) .collect() @@ -295,7 +297,8 @@ mod tests { use openai_protocol::{ common::Function, responses::{ - CodeInterpreterTool, FunctionTool, McpTool, ResponseTool, WebSearchPreviewTool, + CodeInterpreterTool, FunctionTool, ImageGenerationTool, McpTool, ResponseTool, + WebSearchPreviewTool, }, }; use serde_json::json; @@ -525,6 +528,59 @@ mod tests { ); } + #[tokio::test] + async fn test_collect_builtin_routing_image_generation() { + // Image generation is wired through the same hosted-tool MCP plumbing + // as web_search / code_interpreter / file_search. This test proves + // BuiltinToolType::ImageGeneration → ResponseFormat::ImageGenerationCall + // so PRs R6.2/R6.3/R6.4 can rely on the infrastructure. + let mut image_gen_tools = HashMap::new(); + image_gen_tools.insert( + "generate_image".to_string(), + ToolConfig { + response_format: ResponseFormatConfig::ImageGenerationCall, + ..Default::default() + }, + ); + + let config = McpConfig { + servers: vec![McpServerConfig { + name: "image-server".to_string(), + transport: McpTransport::Streamable { + url: "http://localhost:9997/image".to_string(), + token: None, + headers: HashMap::new(), + }, + proxy: None, + required: false, + tools: Some(image_gen_tools), + builtin_type: Some(BuiltinToolType::ImageGeneration), + builtin_tool_name: Some("generate_image".to_string()), + internal: false, + }], + pool: Default::default(), + proxy: None, + warmup: Vec::new(), + inventory: Default::default(), + policy: Default::default(), + }; + + let orchestrator = Arc::new(McpOrchestrator::new(config).await.unwrap()); + + let tools = vec![ResponseTool::ImageGeneration(ImageGenerationTool::default())]; + + let routing = collect_builtin_routing(&orchestrator, Some(&tools)); + + assert_eq!(routing.len(), 1); + assert_eq!(routing[0].builtin_type, BuiltinToolType::ImageGeneration); + assert_eq!(routing[0].server_name, "image-server"); + assert_eq!(routing[0].tool_name, "generate_image"); + assert_eq!( + routing[0].response_format, + ResponseFormat::ImageGenerationCall + ); + } + // ========================================================================= // ensure_request_mcp_client tests // ========================================================================= diff --git a/model_gateway/src/routers/grpc/common/responses/streaming.rs b/model_gateway/src/routers/grpc/common/responses/streaming.rs index b5300b2e1b..7bf6d43d8a 100644 --- a/model_gateway/src/routers/grpc/common/responses/streaming.rs +++ b/model_gateway/src/routers/grpc/common/responses/streaming.rs @@ -8,7 +8,8 @@ use openai_protocol::{ common::{Usage, UsageInfo}, event_types::{ CodeInterpreterCallEvent, ContentPartEvent, FileSearchCallEvent, FunctionCallEvent, - McpEvent, OutputItemEvent, OutputTextEvent, ResponseEvent, WebSearchCallEvent, + ImageGenerationCallEvent, McpEvent, OutputItemEvent, OutputTextEvent, ResponseEvent, + WebSearchCallEvent, }, responses::{ ResponseOutputItem, ResponseStatus, ResponsesRequest, ResponsesResponse, ResponsesUsage, @@ -32,6 +33,7 @@ pub(crate) enum OutputItemType { WebSearchCall, CodeInterpreterCall, FileSearchCall, + ImageGenerationCall, } /// Status of an output item @@ -475,12 +477,18 @@ impl ResponseStreamEventEmitter { ResponseFormat::WebSearchCall => WebSearchCallEvent::IN_PROGRESS, ResponseFormat::CodeInterpreterCall => CodeInterpreterCallEvent::IN_PROGRESS, ResponseFormat::FileSearchCall => FileSearchCallEvent::IN_PROGRESS, + ResponseFormat::ImageGenerationCall => ImageGenerationCallEvent::IN_PROGRESS, ResponseFormat::Passthrough => McpEvent::CALL_IN_PROGRESS, }; self.emit_tool_event(event_type, output_index, item_id) } - /// Emit the searching/interpreting event for builtin tool calls (no-op for passthrough) + /// Emit the searching/interpreting/generating event for builtin tool calls (no-op for passthrough). + /// + /// For `image_generation_call` this emits the `generating` event. The + /// partial-image event is emitted separately via `emit_image_generation_partial_image` + /// because it carries additional payload (the partial b64 bytes) and is + /// optional per the `partial_images` request field. pub fn emit_tool_call_searching( &mut self, output_index: usize, @@ -491,11 +499,46 @@ impl ResponseStreamEventEmitter { ResponseFormat::WebSearchCall => WebSearchCallEvent::SEARCHING, ResponseFormat::CodeInterpreterCall => CodeInterpreterCallEvent::INTERPRETING, ResponseFormat::FileSearchCall => FileSearchCallEvent::SEARCHING, + ResponseFormat::ImageGenerationCall => ImageGenerationCallEvent::GENERATING, ResponseFormat::Passthrough => return None, }; Some(self.emit_tool_event(event_type, output_index, item_id)) } + /// Emit a `response.image_generation_call.partial_image` event. + /// + /// Returns `None` when `response_format` is anything other than + /// [`ResponseFormat::ImageGenerationCall`], mirroring how + /// `emit_tool_call_searching` gates on format. The payload carries the + /// base64-encoded partial image bytes plus a 0-based partial image index. + /// + /// Per-router wiring in R6.2/R6.3/R6.4 is responsible for deciding when + /// to call this and how to source the partial-image bytes. + #[expect( + dead_code, + reason = "partial_image emission is wired by per-router PRs R6.2/R6.3/R6.4" + )] + pub fn emit_image_generation_partial_image( + &mut self, + output_index: usize, + item_id: &str, + response_format: &ResponseFormat, + partial_image_index: u32, + partial_image_b64: &str, + ) -> Option { + if !matches!(response_format, ResponseFormat::ImageGenerationCall) { + return None; + } + Some(json!({ + "type": ImageGenerationCallEvent::PARTIAL_IMAGE, + "sequence_number": self.next_sequence(), + "output_index": output_index, + "item_id": item_id, + "partial_image_index": partial_image_index, + "partial_image_b64": partial_image_b64 + })) + } + /// Emit the appropriate completed event based on response format pub fn emit_tool_call_completed( &mut self, @@ -507,6 +550,7 @@ impl ResponseStreamEventEmitter { ResponseFormat::WebSearchCall => WebSearchCallEvent::COMPLETED, ResponseFormat::CodeInterpreterCall => CodeInterpreterCallEvent::COMPLETED, ResponseFormat::FileSearchCall => FileSearchCallEvent::COMPLETED, + ResponseFormat::ImageGenerationCall => ImageGenerationCallEvent::COMPLETED, ResponseFormat::Passthrough => McpEvent::CALL_COMPLETED, }; self.emit_tool_event(event_type, output_index, item_id) @@ -522,6 +566,7 @@ impl ResponseStreamEventEmitter { Some(ResponseFormat::WebSearchCall) => "web_search_call", Some(ResponseFormat::CodeInterpreterCall) => "code_interpreter_call", Some(ResponseFormat::FileSearchCall) => "file_search_call", + Some(ResponseFormat::ImageGenerationCall) => "image_generation_call", Some(ResponseFormat::Passthrough) => "mcp_call", None => "function_call", } @@ -533,6 +578,7 @@ impl ResponseStreamEventEmitter { Some(ResponseFormat::WebSearchCall) => OutputItemType::WebSearchCall, Some(ResponseFormat::CodeInterpreterCall) => OutputItemType::CodeInterpreterCall, Some(ResponseFormat::FileSearchCall) => OutputItemType::FileSearchCall, + Some(ResponseFormat::ImageGenerationCall) => OutputItemType::ImageGenerationCall, Some(ResponseFormat::Passthrough) => OutputItemType::McpCall, None => OutputItemType::FunctionCall, } @@ -626,6 +672,7 @@ impl ResponseStreamEventEmitter { OutputItemType::WebSearchCall => "ws", OutputItemType::CodeInterpreterCall => "ci", OutputItemType::FileSearchCall => "fs", + OutputItemType::ImageGenerationCall => "ig", }; let id = Self::generate_item_id(id_prefix); diff --git a/model_gateway/src/routers/openai/mcp/tool_loop.rs b/model_gateway/src/routers/openai/mcp/tool_loop.rs index 305843864f..97bbf15a04 100644 --- a/model_gateway/src/routers/openai/mcp/tool_loop.rs +++ b/model_gateway/src/routers/openai/mcp/tool_loop.rs @@ -14,8 +14,8 @@ use axum::http::HeaderMap; use bytes::Bytes; use openai_protocol::{ event_types::{ - is_function_call_type, CodeInterpreterCallEvent, FileSearchCallEvent, ItemType, McpEvent, - OutputItemEvent, WebSearchCallEvent, + is_function_call_type, CodeInterpreterCallEvent, FileSearchCallEvent, + ImageGenerationCallEvent, ItemType, McpEvent, OutputItemEvent, WebSearchCallEvent, }, responses::{generate_id, ResponseInput, ResponseTool, ResponsesRequest}, }; @@ -466,6 +466,11 @@ fn send_tool_call_intermediate_event( ResponseFormat::WebSearchCall => WebSearchCallEvent::SEARCHING, ResponseFormat::CodeInterpreterCall => CodeInterpreterCallEvent::INTERPRETING, ResponseFormat::FileSearchCall => FileSearchCallEvent::SEARCHING, + // stubbed: full event wiring (incl. partial_image payload events) is + // implemented in R6.2 for this router. Emitting `generating` as the + // generic intermediate keeps the shape consistent with the shared + // emitter in `grpc/common/responses/streaming.rs`. + ResponseFormat::ImageGenerationCall => ImageGenerationCallEvent::GENERATING, ResponseFormat::Passthrough => return true, // mcp_call has no intermediate event }; @@ -508,6 +513,7 @@ fn send_tool_call_completion_events( ItemType::WEB_SEARCH_CALL => WebSearchCallEvent::COMPLETED, ItemType::CODE_INTERPRETER_CALL => CodeInterpreterCallEvent::COMPLETED, ItemType::FILE_SEARCH_CALL => FileSearchCallEvent::COMPLETED, + ItemType::IMAGE_GENERATION_CALL => ImageGenerationCallEvent::COMPLETED, _ => McpEvent::CALL_COMPLETED, // Default to mcp_call for mcp_call and unknown types }; @@ -553,6 +559,9 @@ fn stable_streaming_tool_item_id( ResponseFormat::WebSearchCall => normalize_tool_item_id_with_prefix(source_id, "ws_"), ResponseFormat::CodeInterpreterCall => normalize_tool_item_id_with_prefix(source_id, "ci_"), ResponseFormat::FileSearchCall => normalize_tool_item_id_with_prefix(source_id, "fs_"), + // stubbed: `ig_` prefix matches the shared transformer's output item id + // (`to_image_generation_call`). R6.2 wires the full per-router path. + ResponseFormat::ImageGenerationCall => normalize_tool_item_id_with_prefix(source_id, "ig_"), } } @@ -573,7 +582,10 @@ fn non_streaming_tool_item_id_source(item_id: &str, response_format: &ResponseFo ResponseFormat::Passthrough => item_id.to_string(), ResponseFormat::WebSearchCall | ResponseFormat::CodeInterpreterCall - | ResponseFormat::FileSearchCall => item_id + | ResponseFormat::FileSearchCall + // stubbed: image_generation_call shares the fc_/call_ strip behavior + // with the other built-ins. R6.2 wires per-router specifics. + | ResponseFormat::ImageGenerationCall => item_id .strip_prefix("fc_") .or_else(|| item_id.strip_prefix("call_")) .unwrap_or(item_id)