Support compartment

[TingtingZhou7]

Description

Problem

Solution

Changes

Test Plan

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

Summary by CodeRabbit

• New Features
  • Added image generation as a built-in tool with MCP support. The system now handles image generation requests and responses, tracks generation status (InProgress, Generating, Completed, Failed), and preserves output for proper context processing.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: fbecf888-84a3-4319-9e12-03ef324f6b38

📥 Commits

Reviewing files that changed from the base of the PR and between bde1155 and 5830bda.

📒 Files selected for processing (22)

• .gitignore
• crates/mcp/src/core/config.rs
• crates/mcp/src/core/orchestrator.rs
• crates/mcp/src/lib.rs
• crates/mcp/src/transform/mod.rs
• crates/mcp/src/transform/transformer.rs
• crates/mcp/src/transform/types.rs
• crates/protocols/src/event_types.rs
• crates/protocols/src/responses.rs
• model_gateway/src/routers/grpc/common/responses/streaming.rs
• model_gateway/src/routers/grpc/common/responses/utils.rs
• model_gateway/src/routers/grpc/harmony/builder.rs
• model_gateway/src/routers/grpc/regular/responses/common.rs
• model_gateway/src/routers/grpc/regular/responses/non_streaming.rs
• model_gateway/src/routers/grpc/regular/responses/streaming.rs
• model_gateway/src/routers/mcp_utils.rs
• model_gateway/src/routers/mod.rs
• model_gateway/src/routers/openai/mcp/tool_loop.rs
• model_gateway/src/routers/openai/responses/non_streaming.rs
• model_gateway/src/routers/openai/responses/streaming.rs
• model_gateway/src/routers/openai/responses/utils.rs
• model_gateway/src/routers/tool_output_context.rs

---

📝 Walkthrough

Walkthrough

This PR introduces comprehensive support for image generation tool calls across the MCP, protocol, and model gateway layers. It adds new configuration types, response formats, event handlers, and tool execution logic to enable end-to-end image generation functionality.

Changes

Cohort / File(s)	Summary
Configuration & Core Types .gitignore, crates/mcp/src/core/config.rs, crates/mcp/src/core/orchestrator.rs	Added mcp.local.yaml to ignore list, introduced BuiltinToolType::ImageGeneration and ResponseFormatConfig::ImageGenerationCall config variants, and added find_server_config() helper method to orchestrator.
Protocol Definitions crates/protocols/src/event_types.rs, crates/protocols/src/responses.rs	Extended ItemType enum with ImageGenerationCall variant and created new ImageGenerationCallEvent with states (InProgress, Generating, PartialImage, Completed). Added ResponseTool::ImageGeneration, ImageGenerationTool struct, and ResponseOutputItem::ImageGenerationCall with associated ImageGenerationCallStatus enum.
Response Transformation crates/mcp/src/transform/types.rs, crates/mcp/src/transform/transformer.rs, crates/mcp/src/transform/mod.rs, crates/mcp/src/lib.rs	Added ResponseFormat::ImageGenerationCall variant and transformation logic. Implemented extract_image_generation_fallback_text() and is_image_generation_error() helpers for parsing image generation responses. Updated ResponseTransformer::transform() signature to accept &Value instead of &serde_json::Value.
Model Gateway Streaming model_gateway/src/routers/grpc/common/responses/streaming.rs, model_gateway/src/routers/openai/responses/streaming.rs	Added event emission and item-type mapping for ImageGenerationCall format, including ig_ ID prefix and in-progress event injection.
MCP Routing & Tool Execution model_gateway/src/routers/mcp_utils.rs, model_gateway/src/routers/openai/mcp/tool_loop.rs, model_gateway/src/routers/grpc/harmony/builder.rs, model_gateway/src/routers/grpc/common/responses/utils.rs	Extended builtin-tool detection to include ResponseTool::ImageGeneration, added image_generation to tool allowlist, updated ensure_request_mcp_client signature to accept optional request_headers for header forwarding, implemented resolve_forwarded_headers() for MCP transport configuration, and added tool-argument override logic.
Tool Output Processing model_gateway/src/routers/tool_output_context.rs, model_gateway/src/routers/grpc/regular/responses/common.rs, model_gateway/src/routers/grpc/regular/responses/non_streaming.rs, model_gateway/src/routers/grpc/regular/responses/streaming.rs, model_gateway/src/routers/openai/responses/non_streaming.rs, model_gateway/src/routers/openai/responses/utils.rs	Added compact_tool_output_for_model_context() to format image generation outputs, updated ToolLoopState::record_call() signature to accept response_format instead of boolean _success, serialized ResponseTool::ImageGeneration in tool restoration, and threaded response_format through tool execution pipelines.
Module Organization model_gateway/src/routers/mod.rs	Added tool_output_context module declaration.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• refactor(protocol): model ResponseTool as tagged enum to match Responses spec and tighten MCP validation #532 — Both PRs modify ResponseTool variants and MCP routing codepaths in protocols and model_gateway, extending tool support infrastructure.
• fix(openai): unify responses upstream header handling #1029 — Both PRs update ensure_request_mcp_client signature and header propagation for MCP client setup across streaming and non-streaming flows.
• feat(message api): MCP tool use with streaming and non-streaming support #352 — Both PRs modify MCP tool orchestration, connection setup, and tool-loop integration within the model gateway.

Suggested labels

mcp, model-gateway, openai, grpc, protocols, tests

Suggested reviewers

• CatherineSue
• key4ng
• slin1237

Poem

🐰 A hop, skip, and a generate bound,
Image tools now fully aground!
New events stream, routing flows free,
From config to gateway, handlers agree!
Let the generation commence with glee! 🖼️

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces support for a built-in image generation tool across the MCP and model gateway crates. It adds the necessary data structures, response formats, and streaming events to handle image generation calls, including logic to extract base64 image data and metadata from tool results. A significant addition is the compact_tool_output_for_model_context utility, which ensures that large binary image payloads are replaced with a minimal summary before being fed back into the model's conversation history. Furthermore, the PR implements header forwarding for SSE and streamable MCP transports and allows request-level tool overrides to be merged into tool-call arguments. I have no feedback to provide.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5830bda787

ℹ️ 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".

[chatgpt-codex-connector]

Parse SSE-wrapped image tool output before fallback

This fallback runs whenever image_payload_from_wrapped_content returns None, but that helper only handles object-shaped result.content payloads. If the MCP response arrives as SSE text (for example event: message\ndata: {...}), the parser misses it and this branch stores the entire SSE frame as image_generation_call.result, so clients receive protocol framing text instead of the actual image/error payload.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Treat non-wrapper image tool errors as failed in context

This logic only considers result.isError, but tool failures in the main execution path are also emitted as { "error": ... } (for example from execute_tool_resolved on call failures). In those cases is_error is false here, so failed image calls are summarized as status: completed with a success note, which pollutes the tool-loop conversation state and can cause the next model turn to assume success instead of recovering.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Avoid adding both dynamic and static bindings for builtins

When forwarded headers are configured, this block adds a dynamic binding for the builtin server (URL-keyed), but the same builtin type is still passed to ensure_mcp_servers, which also appends the static binding (name-keyed). Because deduping is by server_key, both bindings survive, exposing duplicate tool entries and allowing one route that bypasses forwarded request headers, leading to inconsistent auth behavior depending on which duplicate tool is selected.

Useful? React with 👍 / 👎.