Tingzhou/OpenAI imagegen ci test - #1094
TingtingZhou7 wants to merge 2 commits into
Conversation
github-actions
Bot
added
grpc
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 21 minutes and 27 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (18)
📝 WalkthroughWalkthroughThis pull request adds comprehensive image generation support across the MCP, protocols, and model gateway layers. It introduces new tool types, event types, response formats, transformer logic to extract image payloads from MCP responses, and updated routing infrastructure for both gRPC and OpenAI response paths, along with an end-to-end integration test. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ModelGateway as Model Gateway
participant MCP as MCP Server
participant Transformer
participant Stream as Response Stream
Client->>ModelGateway: Responses API Request<br/>(tools=[image_generation])
ModelGateway->>ModelGateway: Route builtin tool type<br/>to MCP server
ModelGateway->>MCP: Execute image generation<br/>via MCP
MCP-->>ModelGateway: JSON-RPC Response<br/>(wrapped result.content)
ModelGateway->>Transformer: Transform to<br/>ImageGenerationCall format
Transformer->>Transformer: Extract image payload<br/>from wrapped content
Transformer-->>ModelGateway: ImageGenerationCall<br/>ResponseOutputItem
ModelGateway->>Stream: Emit IN_PROGRESS event
ModelGateway->>Stream: Emit GENERATING event
ModelGateway->>Stream: Emit COMPLETED event<br/>with image data
Stream-->>Client: Streaming response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> revert(grpc): drop grpc changes from readded image commit Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> chore: add signed-off empty commit Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> fix(grpc): reject image_generation tool while keeping build exhaustive Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> revert(grpc): drop grpc response/builder changes from branch Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> fix(router): restore literal tool name in compact output context Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> fix(grpc): use ImageGenerationCallEvent for image generation stream events Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com> refactor(mcp): reuse parsed object when extracting image result Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com>
Signed-off-by: TingtingZhou7 <zhoutt96@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a built-in image generation tool within the MCP framework, allowing tool results to be transformed into OpenAI-compatible image generation calls. It introduces new configuration variants, response transformers for image payload extraction, and a tool output compaction utility to summarize results for model context. The PR also updates streaming event types and routing logic. Review feedback identified a need to correctly report failure statuses in image generation calls, recommended replacing fixed sleeps in tests with robust polling, and suggested parameterizing hardcoded tool names in the compaction logic.
| .as_object() | ||
| .or_else(|| parsed_payload.as_ref().and_then(|v| v.as_object())); | ||
|
|
||
| let status = ImageGenerationCallStatus::Completed; |
There was a problem hiding this comment.
The status of the image generation call is always set to Completed, even when an error is detected. This can lead to incorrect reporting of tool call outcomes. The status should be set to Failed when is_image_generation_error returns true. Additionally, if this outcome is recorded for a worker request, ensure that the record_outcome function is called with the HTTP status code, not a boolean, to allow the internal logic to determine if it represents a failure.
| let status = ImageGenerationCallStatus::Completed; | |
| let status = if is_image_generation_error(result) { | |
| ImageGenerationCallStatus::Failed | |
| } else { | |
| ImageGenerationCallStatus::Completed | |
| }; |
References
- When recording the outcome of a worker request, ensure that the record_outcome function is called with the HTTP status code, not a boolean indicating success or failure. The worker's internal logic will determine if the status code represents a circuit breaker failure based on its configured retryable status codes.
| """Test that image_generation produces image_generation_call output.""" | ||
| gateway, client = gateway_with_mcp_config | ||
|
|
||
| time.sleep(2) |
There was a problem hiding this comment.
Using a fixed time.sleep() can lead to flaky tests, as the required wait time might vary. If this sleep is intended to wait for the gateway to be ready, consider moving this logic into the gateway_with_mcp_config fixture and implementing a more robust waiting mechanism, such as polling a health check endpoint. Each process should have its own individual startup timeout rather than using a shared deadline.
References
- When waiting for multiple concurrent processes to become healthy, each process should have its own individual startup timeout, rather than using a shared deadline for all.
| "Successfully generated the image".to_string() | ||
| }; | ||
| let summary = json!({ | ||
| "tool": "generate_image", |
There was a problem hiding this comment.
The tool name in the summary is hardcoded to "generate_image". This might be inaccurate if the actual tool name is different. Consider passing the tool name as an argument to compact_tool_output_for_model_context and using it here to make the summary more accurate.
For example, you could change the function signature to:
pub fn compact_tool_output_for_model_context(
response_format: &ResponseFormat,
output: &Value,
tool_name: &str,
) -> String
And then use tool_name here.
| "tool": "generate_image", | |
| "tool": tool_name, |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 01365fbc34
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| .as_object() | ||
| .or_else(|| parsed_payload.as_ref().and_then(|v| v.as_object())); | ||
|
|
||
| let status = ImageGenerationCallStatus::Completed; |
There was a problem hiding this comment.
to_image_generation_call hard-codes status to completed, so even error payloads (for example responses with result.isError=true and error text) are emitted as successful image_generation_call items. Clients that rely on status to drive retry/error handling will mis-handle real tool failures as successes.
Useful? React with 👍 / 👎.
| let is_error = is_image_generation_error(output); | ||
| let note = if is_error { | ||
| extract_image_generation_fallback_text(output).unwrap_or_default() | ||
| } else { | ||
| "Successfully generated the image".to_string() | ||
| }; | ||
| let summary = json!({ | ||
| "tool": "generate_image", | ||
| "status": if is_error { "failed" } else { "completed" }, |
There was a problem hiding this comment.
This compaction path infers failure only from result.isError; when image tool execution fails with other common shapes (e.g. {"error": ...}), it still writes a completed summary with a success note. Because this summary is fed into conversation history, the next model turn is given incorrect success context and can generate a wrong follow-up response.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
model_gateway/src/routers/grpc/harmony/builder.rs (1)
418-435:⚠️ Potential issue | 🟠 MajorDon’t classify
image_generationas a custom Harmony tool.Line 428 adds
"image_generation"to thetool_typeslist, buthas_custom_tools()still only treatsweb_search_preview,code_interpreter, andcontaineras builtins. That makes image-generation-only requests setwith_custom_tools = true, even thoughbuild_developer_message_from_responses()only emits descriptions forResponseTool::Function(_). The Responses Harmony path will therefore take the custom-tool prompt branch for a builtin and can emit an empty developer message.Suggested fix
- let tool_types: Vec<&str> = request - .tools - .as_ref() - .map(|tools| { - tools - .iter() - .map(|tool| match tool { - ResponseTool::Function(_) => "function", - ResponseTool::WebSearchPreview(_) => "web_search_preview", - ResponseTool::CodeInterpreter(_) => "code_interpreter", - ResponseTool::ImageGeneration(_) => "image_generation", - ResponseTool::Mcp(_) => "mcp", - }) - .collect() - }) - .unwrap_or_default(); - - let with_custom_tools = has_custom_tools(&tool_types); + let with_custom_tools = request.tools.as_ref().is_some_and(|tools| { + tools + .iter() + .any(|tool| matches!(tool, ResponseTool::Function(_) | ResponseTool::Mcp(_))) + });🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/grpc/harmony/builder.rs` around lines 418 - 435, tool_types currently includes "image_generation" which causes with_custom_tools (called via has_custom_tools) to return true even though image generation is a builtin; update the logic so image_generation is treated as a builtin. Fix by modifying has_custom_tools to exclude "image_generation" from the custom-tools set (or add "image_generation" to its builtin list), or alternatively filter out "image_generation" from tool_types before calling has_custom_tools; ensure this aligns with ResponseTool::ImageGeneration and with build_developer_message_from_responses so image-generation-only requests do not flip the custom-tool prompt branch.model_gateway/src/routers/mcp_utils.rs (1)
133-158:⚠️ Potential issue | 🟠 MajorHandle
ResponseTool::FileSearch(_)in both builtin-routing helpers.The updated docstrings now list
file_search, but both match statements still skipResponseTool::FileSearch(_). As a result, file-search builtins never make it into builtin routing orensure_request_mcp_client(), so a configured MCP server for{type: "file_search"}still won't be attached on this path.Also applies to: 192-201
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/mcp_utils.rs` around lines 133 - 158, The builtin routing helpers currently ignore ResponseTool::FileSearch and so file_search never gets routed; update collect_builtin_routing to include a match arm mapping ResponseTool::FileSearch(_) => BuiltinToolType::FileSearch and likewise update the other helper (ensure_request_mcp_client / the second match block that mirrors lines ~192-201) to handle ResponseTool::FileSearch(_) so configured MCP servers for type "file_search" are considered; make sure the BuiltinToolType enum variant FileSearch is used in both places to mirror the other builtins.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/routers/openai/mcp/tool_loop.rs`:
- Around line 268-313: request_tool_overrides currently serializes the entire
ResponseTool::ImageGeneration object and can merge wrapper or deferred keys
(like type or per-option image settings) into MCP args; change
request_tool_overrides to only extract and return a JSON object containing the
explicitly supported image-generation override keys (e.g., "model" and
"revised_prompt")—dropping nulls—and nothing else so
apply_request_tool_overrides will only insert those allowed keys; keep the
null-filtering behavior and ensure behavior is consistent with
sanitize_builtin_tool_arguments for ResponseFormat::ImageGenerationCall and
return None when no supported overrides are present.
---
Outside diff comments:
In `@model_gateway/src/routers/grpc/harmony/builder.rs`:
- Around line 418-435: tool_types currently includes "image_generation" which
causes with_custom_tools (called via has_custom_tools) to return true even
though image generation is a builtin; update the logic so image_generation is
treated as a builtin. Fix by modifying has_custom_tools to exclude
"image_generation" from the custom-tools set (or add "image_generation" to its
builtin list), or alternatively filter out "image_generation" from tool_types
before calling has_custom_tools; ensure this aligns with
ResponseTool::ImageGeneration and with build_developer_message_from_responses so
image-generation-only requests do not flip the custom-tool prompt branch.
In `@model_gateway/src/routers/mcp_utils.rs`:
- Around line 133-158: The builtin routing helpers currently ignore
ResponseTool::FileSearch and so file_search never gets routed; update
collect_builtin_routing to include a match arm mapping
ResponseTool::FileSearch(_) => BuiltinToolType::FileSearch and likewise update
the other helper (ensure_request_mcp_client / the second match block that
mirrors lines ~192-201) to handle ResponseTool::FileSearch(_) so configured MCP
servers for type "file_search" are considered; make sure the BuiltinToolType
enum variant FileSearch is used in both places to mirror the other builtins.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: eb8217d3-6ed8-4824-8d64-b2a6ce6b1d30
📒 Files selected for processing (18)
.gitignorecrates/mcp/src/core/config.rscrates/mcp/src/lib.rscrates/mcp/src/transform/mod.rscrates/mcp/src/transform/transformer.rscrates/mcp/src/transform/types.rscrates/protocols/src/event_types.rscrates/protocols/src/responses.rse2e_test/responses/test_builtin_tools.pymodel_gateway/src/routers/grpc/common/responses/streaming.rsmodel_gateway/src/routers/grpc/common/responses/utils.rsmodel_gateway/src/routers/grpc/harmony/builder.rsmodel_gateway/src/routers/mcp_utils.rsmodel_gateway/src/routers/mod.rsmodel_gateway/src/routers/openai/mcp/tool_loop.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/routers/openai/responses/utils.rsmodel_gateway/src/routers/tool_output_context.rs
| fn request_tool_overrides( | ||
| response_format: &ResponseFormat, | ||
| original_body: &ResponsesRequest, | ||
| ) -> Option<Value> { | ||
| if !matches!(response_format, ResponseFormat::ImageGenerationCall) { | ||
| return None; | ||
| } | ||
|
|
||
| // Read request-defined tools and find the image_generation config. | ||
| let tools = original_body.tools.as_ref()?; | ||
|
|
||
| tools.iter().find_map(|tool| { | ||
| // Serialize image tool config into a JSON object for merge. | ||
| let mut serialized = match tool { | ||
| ResponseTool::ImageGeneration(image_tool) => match to_value(image_tool).ok()? { | ||
| Value::Object(obj) => obj, | ||
| _ => return None, | ||
| }, | ||
| _ => return None, | ||
| }; | ||
| // Drop nulls so absent fields do not overwrite generated call arguments. | ||
| 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, | ||
| ) { | ||
| if let (Some(overrides), Some(args_obj)) = ( | ||
| request_tool_overrides(response_format, original_body), | ||
| arguments.as_object_mut(), | ||
| ) { | ||
| let Some(override_obj) = overrides.as_object() else { | ||
| return; | ||
| }; | ||
| for (k, v) in override_obj { | ||
| args_obj.insert(k.clone(), v.clone()); | ||
| } | ||
| } |
There was a problem hiding this comment.
Filter image-generation overrides before merging them into MCP args.
request_tool_overrides() serializes the entire ResponseTool::ImageGeneration object and merges every non-null field into the execution payload. That reintroduces wrapper keys like type and any currently-deferred image options into the MCP call, which can break tools that only accept the sanitized image-generation argument subset. Please restrict this merge to the explicitly supported keys instead of cloning the whole tool object.
Based on learnings, sanitize_builtin_tool_arguments in model_gateway/src/routers/openai/mcp/tool_loop.rs intentionally keeps only model and revised_prompt for ResponseFormat::ImageGenerationCall; per-option overrides/defaults are deferred.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/src/routers/openai/mcp/tool_loop.rs` around lines 268 - 313,
request_tool_overrides currently serializes the entire
ResponseTool::ImageGeneration object and can merge wrapper or deferred keys
(like type or per-option image settings) into MCP args; change
request_tool_overrides to only extract and return a JSON object containing the
explicitly supported image-generation override keys (e.g., "model" and
"revised_prompt")—dropping nulls—and nothing else so
apply_request_tool_overrides will only insert those allowed keys; keep the
null-filtering behavior and ensure behavior is consistent with
sanitize_builtin_tool_arguments for ResponseFormat::ImageGenerationCall and
return None when no supported overrides are present.
1 participant
Description
Problem
Solution
Changes
Test Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit