resolve comments from simo

Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com>

Author: TingtingZhou7

.gitignore

@@ -160,3 +160,6 @@ clients/java/target/
 clients/java/build.sbt
 clients/java/gradle.properties
 openapitools.json
+
+# Local MCP config
+mcp.local.yaml

crates/mcp/src/transform/transformer.rs

@@ -41,12 +41,8 @@ impl ResponseTransformer {
     // (derived from builtin_type), not payload key heuristics.
     fn image_payload_from_wrapped_content(
         result: &Value,
-        format: &ResponseFormat,
+        _format: &ResponseFormat,
     ) -> Option<Value> {
-        if !matches!(format, ResponseFormat::ImageGenerationCall) {
-            return None;
-        }
-
         // Already a direct object payload.
         if result
             .as_object()

crates/protocols/src/responses.rs

@@ -86,12 +86,19 @@ pub struct CodeInterpreterTool {
 
 /// Built-in image generation tool.
 ///
-/// The upstream API may add optional fields over time. We keep this tool shape
-/// open by preserving unknown key/value pairs.
+/// Known fields are typed for validation; unknown fields are preserved for
+/// forward compatibility via `extra`.
 #[derive(Debug, Clone, Deserialize, Serialize, Default, schemars::JsonSchema)]
 pub struct ImageGenerationTool {
+    pub size: Option<String>,
+    pub quality: Option<String>,
+    pub background: Option<String>,
+    pub output_format: Option<String>,
+    pub output_compression: Option<u32>,
+    pub moderation: Option<String>,
+    pub model: Option<String>,
     #[serde(flatten)]
-    pub options: HashMap<String, Value>,
+    pub extra: HashMap<String, Value>,
 }
 
 /// `require_approval` values.
@@ -1506,14 +1513,16 @@ mod tests {
         let tool: ResponseTool = serde_json::from_value(json!({
             "type": "image_generation",
             "size": "1024x1024",
-            "quality": "high"
+            "quality": "high",
+            "custom_flag": true
         }))
         .expect("image_generation tool should deserialize");
 
         match tool {
             ResponseTool::ImageGeneration(t) => {
-                assert_eq!(t.options.get("size"), Some(&json!("1024x1024")));
-                assert_eq!(t.options.get("quality"), Some(&json!("high")));
+                assert_eq!(t.size.as_deref(), Some("1024x1024"));
+                assert_eq!(t.quality.as_deref(), Some("high"));
+                assert_eq!(t.extra.get("custom_flag"), Some(&json!(true)));
             }
             other => panic!("expected image_generation tool, got {other:?}"),
         }

mcp.local.yaml

@@ -7,3 +7,11 @@ servers:
     tools:
       generate_image:
         response_format: image_generation_call
+  - name: search_web
+    protocol: streamable
+    url: http://localhost:8080/mcp
+    builtin_type: web_search_preview
+    builtin_tool_name: search_web
+    tools:
+      search_web:
+        response_format: web_search_call

model_gateway/src/routers/openai/mcp/tool_loop.rs

@@ -17,7 +17,7 @@ use openai_protocol::{
         is_function_call_type, CodeInterpreterCallEvent, FileSearchCallEvent,
         ImageGenerationCallEvent, ItemType, McpEvent, OutputItemEvent, WebSearchCallEvent,
     },
-    responses::{generate_id, ResponseInput, ResponsesRequest},
+    responses::{generate_id, ResponseInput, ResponseTool, ResponsesRequest},
 };
 use serde_json::{json, to_value, Value};
 use smg_mcp::{
@@ -35,8 +35,8 @@ use crate::{
     },
 };
 
-// TODO: Temporary hardcoded image model; replace with configurable model routing later.
-const IMAGE_MODEL: &str = "openai.gpt-image-1.5";
+// Only image_generation is enabled for request-level argument overrides for now.
+const ALLOWED_BUILTIN_TOOLS_FOR_OVERRIDE: &[&str] = &["image_generation"];
 
 /// State for tracking multi-turn tool calling loop
 pub(crate) struct ToolLoopState {
@@ -103,6 +103,7 @@ pub(crate) async fn execute_streaming_tool_calls(
     state: &mut ToolLoopState,
     sequence_number: &mut u64,
     model_id: &str,
+    original_body: &ResponsesRequest,
 ) -> bool {
     for call in pending_calls {
         if call.name.is_empty() {
@@ -166,13 +167,12 @@ pub(crate) async fn execute_streaming_tool_calls(
                 continue;
             }
         };
-        sanitize_builtin_tool_arguments(&response_format, &mut arguments);
-
+        apply_request_tool_overrides(&response_format, original_body, &mut arguments);
         if !send_tool_call_intermediate_event(tx, &call, &response_format, sequence_number) {
             return false;
         }
 
-        debug!("Calling MCP tool '{}' with args: {}", call.name, args_str);
+        debug!("Calling MCP tool '{}' with args: {}", call.name, arguments);
         let tool_output = session
             .execute_tool(ToolExecutionInput {
                 call_id: call.call_id.clone(),
@@ -261,6 +261,82 @@ pub(crate) fn prepare_mcp_tools_as_functions(payload: &mut Value, session: &McpT
     }
 }
 
+fn builtin_tool_type_for_response_format(response_format: &ResponseFormat) -> Option<&'static str> {
+    match response_format {
+        ResponseFormat::ImageGenerationCall => Some("image_generation"),
+        ResponseFormat::WebSearchCall => Some("web_search_preview"),
+        ResponseFormat::CodeInterpreterCall => Some("code_interpreter"),
+        ResponseFormat::FileSearchCall | ResponseFormat::Passthrough => None,
+    }
+}
+
+fn request_tool_overrides(
+    response_format: &ResponseFormat,
+    original_body: &ResponsesRequest,
+) -> Option<Value> {
+    let builtin_tool_type = builtin_tool_type_for_response_format(response_format)?;
+    if !ALLOWED_BUILTIN_TOOLS_FOR_OVERRIDE.contains(&builtin_tool_type) {
+        return None;
+    }
+
+    let tools = original_body.tools.as_ref()?;
+
+    tools.iter().find_map(|tool| {
+        let mut serialized = match (builtin_tool_type, tool) {
+            ("image_generation", ResponseTool::ImageGeneration(image_tool)) => {
+                match to_value(image_tool).ok()? {
+                    Value::Object(obj) => obj,
+                    _ => return None,
+                }
+            }
+            ("web_search_preview", ResponseTool::WebSearchPreview(web_search_tool)) => {
+                match to_value(web_search_tool).ok()? {
+                    Value::Object(obj) => obj,
+                    _ => return None,
+                }
+            }
+            ("code_interpreter", ResponseTool::CodeInterpreter(code_interpreter_tool)) => {
+                match to_value(code_interpreter_tool).ok()? {
+                    Value::Object(obj) => obj,
+                    _ => return None,
+                }
+            }
+            _ => {
+                return None;
+            }
+        };
+        serialized.retain(|_, v| !v.is_null());
+        if serialized.is_empty() {
+            None
+        } else {
+            Some(Value::Object(serialized))
+        }
+    })
+}
+
+fn apply_request_tool_overrides(
+    response_format: &ResponseFormat,
+    original_body: &ResponsesRequest,
+    arguments: &mut Value,
+) {
+    let Some(overrides) = request_tool_overrides(response_format, original_body) else {
+        return;
+    };
+    let Some(override_obj) = overrides.as_object() else {
+        return;
+    };
+
+    if !arguments.is_object() {
+        *arguments = json!({});
+    }
+
+    if let Some(args_obj) = arguments.as_object_mut() {
+        for (k, v) in override_obj {
+            args_obj.insert(k.clone(), v.clone());
+        }
+    }
+}
+
 /// Build a resume payload with conversation history
 pub(crate) fn build_resume_payload(
     base_payload: &Value,
@@ -683,11 +759,10 @@ pub(crate) async fn execute_tool_loop(
                     continue;
                 }
             };
-            sanitize_builtin_tool_arguments(&response_format, &mut arguments);
-
+            apply_request_tool_overrides(&response_format, original_body, &mut arguments);
             debug!(
                 "Calling MCP tool '{}' with args: {}",
-                call.name, call.arguments
+                call.name, arguments
             );
             let tool_output = session
                 .execute_tool(ToolExecutionInput {
@@ -849,37 +924,6 @@ fn build_mcp_call_item(
     })
 }
 
-// TODO: Extend this with per-tool argument override rules as we add more
-// builtin tool formats that need request-time normalization.
-fn sanitize_builtin_tool_arguments(response_format: &ResponseFormat, arguments: &mut Value) {
-    match response_format {
-        ResponseFormat::ImageGenerationCall => {
-            // Current fallback behavior: we intentionally narrow image arguments
-            // to `model` + `revised_prompt` for compatibility. This drops extra
-            // image options for now; we'll later pass either explicit user-provided
-            // values or safe defaults per option instead of truncating.
-            let revised_prompt = arguments
-                .as_object()
-                .and_then(|obj| {
-                    obj.get("revised_prompt")
-                        .and_then(|v| v.as_str())
-                        .or_else(|| obj.get("prompt").and_then(|v| v.as_str()))
-                })
-                .unwrap_or("")
-                .to_string();
-
-            *arguments = json!({
-                "model": IMAGE_MODEL,
-                "revised_prompt": revised_prompt
-            });
-        }
-        ResponseFormat::WebSearchCall
-        | ResponseFormat::CodeInterpreterCall
-        | ResponseFormat::FileSearchCall
-        | ResponseFormat::Passthrough => {}
-    }
-}
-
 /// Build a transformed output item using ResponseTransformer
 ///
 /// Converts the output using the tool's response_format to the correctly-typed

model_gateway/src/routers/openai/responses/streaming.rs

@@ -968,6 +968,7 @@ pub(super) fn handle_streaming_with_tool_interception(
                 &mut state,
                 &mut sequence_number,
                 &original_request.model,
+                &original_request,
             )
             .await
             {