refactor(protocol): model ResponseTool as tagged enum to match Responses spec and tighten MCP validation - #532
Conversation
📝 WalkthroughWalkthroughThis PR refactors the ResponseTool representation from a struct with a type discriminator field to a strongly-typed enum with variants (Function, Mcp, WebSearchPreview, CodeInterpreter), removing the ResponseToolType enum. The change propagates through MCP bridging, tool extraction utilities, and streaming processors, replacing type-field-based pattern matching with enum-variant matching and removing precomputed MCP tool name sets in favor of session-based exposure checks. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 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 |
…ses spec and tighten MCP validation Signed-off-by: Ziwen Zhao <zzw.mose@gmail.com>
5ec5d4d to
5843ebb
Compare
Summary of ChangesHello @zhaowenzi, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the internal representation of Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request refactors the ResponseTool enum to use type-based discrimination instead of a type field, introducing dedicated structs for FunctionTool, WebSearchPreviewTool, CodeInterpreterTool, and McpTool. The code changes remove ResponseToolType enum. The extract_tools_from_response_tools function is updated to extract only function tools. The review comments identify a Server-Side Request Forgery (SSRF) vulnerability in the ensure_request_mcp_client function due to the lack of validation of the user-provided server_url in the McpTool configuration, which could allow an attacker to point to internal services or cloud metadata endpoints. The validation logic for ResponseTool::Mcp only verifies the server_label format but lacks any validation for the server_url field.
There was a problem hiding this comment.
Actionable comments posted: 2
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/responses/streaming.rs (1)
244-268:⚠️ Potential issue | 🟠 MajorEnsure
max_tool_callscounts function tool calls when MCP is enabled.
total_calls_afteronly includes MCP calls, butmax_tool_callsis a global limit. If the model returns MCP + function tool calls in the same iteration, the limit can be exceeded without emittingincomplete_details. Consider counting both types (or explicitly document MCP-only semantics).🛠️ Suggested fix
- let total_calls_after = mcp_tracking.total_calls() + mcp_tool_calls.len(); + let total_calls_after = mcp_tracking.total_calls() + + mcp_tool_calls.len() + + function_tool_calls.len();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/grpc/harmony/responses/streaming.rs` around lines 244 - 268, The check for exceeding max_tool_calls only adds MCP calls; include function tool calls as well: compute total_calls_after using mcp_tracking.total_calls() + mcp_tool_calls.len() + function_tool_calls.len() (or otherwise sum mcp_tool_calls and function_tool_calls) and compare that sum to effective_limit; update the warn() fields (new_calls/total_after) to reflect the combined new calls and resultant total, and ensure any emitted incomplete_details path uses this combined count logic (refer to mcp_tool_calls, function_tool_calls, total_calls_after, effective_limit, mcp_tracking.total_calls(), and max_tool_calls).mcp/src/responses_bridge.rs (1)
76-99:⚠️ Potential issue | 🟡 MinorUpdate the MCP tool docstring to reflect function-tool construction.
The comment still says MCP tools are represented as
type: "mcp", but the implementation now returnsResponseTool::Function. If this is intentional, update the doc to avoid misleading callers; otherwise, restore the MCP variant in the builder.📝 Suggested doc update
-/// Build Responses API MCP tools from MCP tool entries. -/// -/// These tools are exposed in Responses requests where MCP tools are represented -/// as `{"type": "mcp", ...}` tool entries. +/// Build Responses API function tools from MCP tool entries. +/// +/// These tools are exposed in Responses requests as `{"type": "function", ...}` tool entries.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@mcp/src/responses_bridge.rs` around lines 76 - 99, The docstring for build_response_tools / build_response_tools_with_names is stale: the code constructs ResponseTool::Function (via FunctionTool / Function and using resolved_name_for_entry, entry.tool.input_schema, etc.) but the comment still says MCP tools are represented as `{"type": "mcp", ...}`. Fix by updating the top-of-function comments to state that MCP tool entries are exposed as Responses Function tools (ResponseTool::Function) with name, description, and parameters derived from resolved_name_for_entry, entry.tool.description, and entry.tool.input_schema; alternatively, if the original MCP variant behavior is required, change the mapper to construct the MCP ResponseTool variant instead of ResponseTool::Function. Ensure references to build_response_tools and build_response_tools_with_names remain accurate.
🤖 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/grpc/harmony/responses/non_streaming.rs`:
- Around line 170-176: The partition logic using
session.has_exposed_tool(&tc.function.name) can misclassify user function tools
that share names with MCP-exposed tools; add a preflight check before the
partition to detect any name collisions between the MCP-exposed tool names (from
session.exposed tools, used by session.has_exposed_tool) and the incoming
tool_calls' tc.function.name values, and if any overlaps are found reject the
request with a clear error (or apply an explicit MCP namespace rule) instead of
proceeding to let partition(...) separate them; update the code around
tool_calls, the partition call, and session.has_exposed_tool usage to enforce
this guard.
In `@model_gateway/src/routers/openai/responses/utils.rs`:
- Around line 211-233: The match arm in response_tool_to_value currently only
returns a Value for ResponseTool::Mcp when mcp.server_url.is_some(), which drops
static MCP tools that only have server_label; change the logic in
response_tool_to_value so ResponseTool::Mcp always builds and returns the object
(including type "mcp", server_label, and optional fields server_url,
server_description, require_approval, allowed_tools) regardless of whether
mcp.server_url is Some or None; keep using insert_optional_value for optional
fields and the same allowed_tools handling so static MCP entries are preserved.
---
Outside diff comments:
In `@mcp/src/responses_bridge.rs`:
- Around line 76-99: The docstring for build_response_tools /
build_response_tools_with_names is stale: the code constructs
ResponseTool::Function (via FunctionTool / Function and using
resolved_name_for_entry, entry.tool.input_schema, etc.) but the comment still
says MCP tools are represented as `{"type": "mcp", ...}`. Fix by updating the
top-of-function comments to state that MCP tool entries are exposed as Responses
Function tools (ResponseTool::Function) with name, description, and parameters
derived from resolved_name_for_entry, entry.tool.description, and
entry.tool.input_schema; alternatively, if the original MCP variant behavior is
required, change the mapper to construct the MCP ResponseTool variant instead of
ResponseTool::Function. Ensure references to build_response_tools and
build_response_tools_with_names remain accurate.
In `@model_gateway/src/routers/grpc/harmony/responses/streaming.rs`:
- Around line 244-268: The check for exceeding max_tool_calls only adds MCP
calls; include function tool calls as well: compute total_calls_after using
mcp_tracking.total_calls() + mcp_tool_calls.len() + function_tool_calls.len()
(or otherwise sum mcp_tool_calls and function_tool_calls) and compare that sum
to effective_limit; update the warn() fields (new_calls/total_after) to reflect
the combined new calls and resultant total, and ensure any emitted
incomplete_details path uses this combined count logic (refer to mcp_tool_calls,
function_tool_calls, total_calls_after, effective_limit,
mcp_tracking.total_calls(), and max_tool_calls).
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (15)
mcp/src/responses_bridge.rsmodel_gateway/src/routers/grpc/common/responses/utils.rsmodel_gateway/src/routers/grpc/harmony/builder.rsmodel_gateway/src/routers/grpc/harmony/responses/common.rsmodel_gateway/src/routers/grpc/harmony/responses/non_streaming.rsmodel_gateway/src/routers/grpc/harmony/responses/streaming.rsmodel_gateway/src/routers/grpc/harmony/stages/preparation.rsmodel_gateway/src/routers/grpc/harmony/streaming.rsmodel_gateway/src/routers/grpc/regular/responses/conversions.rsmodel_gateway/src/routers/mcp_utils.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/routers/openai/responses/utils.rsmodel_gateway/tests/api/responses_api_test.rsmodel_gateway/tests/spec/responses.rsprotocols/src/responses.rs
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (3)
model_gateway/src/routers/openai/responses/utils.rs (1)
211-230:⚠️ Potential issue | 🟠 MajorRestore MCP tools even without
server_url.The guard on
mcp.server_url.is_some()still drops MCP tools that only provideserver_label, which breaks mirroring of original request tools.🔧 Suggested fix
- ResponseTool::Mcp(mcp) if mcp.server_url.is_some() => { + ResponseTool::Mcp(mcp) => {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/openai/responses/utils.rs` around lines 211 - 230, The match arm currently only handles ResponseTool::Mcp when mcp.server_url.is_some(), which drops MCP tools that only have server_label; change the pattern to match all ResponseTool::Mcp(mcp) (remove the server_url.is_some() guard) so the block always builds an object and uses insert_optional_value(&mut m, "server_url", mcp.server_url.as_ref()) to include server_url only when present; keep the existing inserts for server_label, server_description, require_approval, and allowed_tools as-is.model_gateway/src/routers/grpc/harmony/responses/streaming.rs (1)
244-248: MCP/function tool name collision risk (same as non-streaming path).Same
session.has_exposed_toolpartition logic asnon_streaming.rsLine 175.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/grpc/harmony/responses/streaming.rs` around lines 244 - 248, The partitioning currently uses session.has_exposed_tool(&tc.function.name) which can collide on duplicate names; update the partition logic in the streaming path (where tool_calls is split into mcp_tool_calls and function_tool_calls) to use a unique identifier or provenance check instead of raw name—e.g., use tc.function.id or a fully-qualified name/namespace field (or compare tc.function.origin/mapping) when calling session.has_exposed_tool (or add a new session method that accepts the function id) so streaming and non-streaming paths disambiguate tools by identity rather than name.model_gateway/src/routers/grpc/harmony/responses/non_streaming.rs (1)
170-175: MCP/function tool name collision risk persists with session-based partitioning.The partition at Line 175 uses
session.has_exposed_tool(&tc.function.name)as the sole discriminator. If a user-defined function tool happens to share a name with an MCP-exposed tool, it will be misclassified as MCP and executed through the MCP path. This is a correctness/security concern.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/grpc/harmony/responses/non_streaming.rs` around lines 170 - 175, The current partition uses only session.has_exposed_tool(&tc.function.name) which misclassifies user tools that share names with MCP tools; change the partition predicate in the tool_calls -> partition(...) logic to require both that the name is exposed AND that the tool call actually originates from the MCP (e.g., check an explicit origin/type field on the ToolCall/Function struct such as tc.function.origin or tc.metadata indicating MCP), falling back to comparing a unique tool id/namespace rather than name; update references around mcp_tool_calls and function_tool_calls accordingly and add tests for a same-name user tool vs MCP tool collision.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@mcp/src/responses_bridge.rs`:
- Around line 91-99: The docstring for build_response_tools_with_names is
incorrect: the function converts MCP tool entries into Responses API
function-type tools (constructed as ResponseTool::Function with FunctionTool and
Function) rather than producing {"type": "mcp", ...} entries; update the
docstring to state that MCP tool entries are transformed into function tools
(serialize as {"type": "function", ...}), mentioning the use of
ResponseTool::Function, FunctionTool and resolved_name_for_entry to make the
behavior clear.
In `@model_gateway/src/routers/openai/responses/utils.rs`:
- Around line 231-232: ResponseTool::WebSearchPreview and
ResponseTool::CodeInterpreter currently call serde_json::to_value(tool) which
produces adjacent-tagged JSON; change both branches to manually construct an
object with an explicit "type" field (e.g. "web_search_preview" and
"code_interpreter") and then merge the tool's inner serialized fields into that
object so the output shape matches the MCP branch. Locate the
ResponseTool::WebSearchPreview and ResponseTool::CodeInterpreter arms and
replace the serde_json::to_value(tool).ok() usage with code that creates a
serde_json::Map, inserts ("type",
serde_json::Value::String("web_search_preview"/"code_interpreter")), serializes
the inner data (via serde_json::to_value on the inner struct) and extends the
map with those key/value pairs, then returns
serde_json::Value::Object(map).ok().
---
Duplicate comments:
In `@model_gateway/src/routers/grpc/harmony/responses/non_streaming.rs`:
- Around line 170-175: The current partition uses only
session.has_exposed_tool(&tc.function.name) which misclassifies user tools that
share names with MCP tools; change the partition predicate in the tool_calls ->
partition(...) logic to require both that the name is exposed AND that the tool
call actually originates from the MCP (e.g., check an explicit origin/type field
on the ToolCall/Function struct such as tc.function.origin or tc.metadata
indicating MCP), falling back to comparing a unique tool id/namespace rather
than name; update references around mcp_tool_calls and function_tool_calls
accordingly and add tests for a same-name user tool vs MCP tool collision.
In `@model_gateway/src/routers/grpc/harmony/responses/streaming.rs`:
- Around line 244-248: The partitioning currently uses
session.has_exposed_tool(&tc.function.name) which can collide on duplicate
names; update the partition logic in the streaming path (where tool_calls is
split into mcp_tool_calls and function_tool_calls) to use a unique identifier or
provenance check instead of raw name—e.g., use tc.function.id or a
fully-qualified name/namespace field (or compare tc.function.origin/mapping)
when calling session.has_exposed_tool (or add a new session method that accepts
the function id) so streaming and non-streaming paths disambiguate tools by
identity rather than name.
In `@model_gateway/src/routers/openai/responses/utils.rs`:
- Around line 211-230: The match arm currently only handles ResponseTool::Mcp
when mcp.server_url.is_some(), which drops MCP tools that only have
server_label; change the pattern to match all ResponseTool::Mcp(mcp) (remove the
server_url.is_some() guard) so the block always builds an object and uses
insert_optional_value(&mut m, "server_url", mcp.server_url.as_ref()) to include
server_url only when present; keep the existing inserts for server_label,
server_description, require_approval, and allowed_tools as-is.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (15)
mcp/src/responses_bridge.rsmodel_gateway/src/routers/grpc/common/responses/utils.rsmodel_gateway/src/routers/grpc/harmony/builder.rsmodel_gateway/src/routers/grpc/harmony/responses/common.rsmodel_gateway/src/routers/grpc/harmony/responses/non_streaming.rsmodel_gateway/src/routers/grpc/harmony/responses/streaming.rsmodel_gateway/src/routers/grpc/harmony/stages/preparation.rsmodel_gateway/src/routers/grpc/harmony/streaming.rsmodel_gateway/src/routers/grpc/regular/responses/conversions.rsmodel_gateway/src/routers/mcp_utils.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/routers/openai/responses/utils.rsmodel_gateway/tests/api/responses_api_test.rsmodel_gateway/tests/spec/responses.rsprotocols/src/responses.rs
…ses spec and tighten MCP validation Signed-off-by: Ziwen Zhao <zzw.mose@gmail.com>
…ses spec and tighten MCP validation Signed-off-by: Ziwen Zhao <zzw.mose@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/grpc/harmony/streaming.rs`:
- Around line 734-741: The code does two sequential lookups via
McpToolSession::has_exposed_tool and then tool_response_format; collapse this to
a single lookup to avoid the redundant map access by adding a new method like
McpToolSession::try_tool_response_format(name: &str) -> Option<ResponseFormat>
(or inline the single lookup using exposed_name_map.get(name).map(|entry|
entry.tool_response_format().unwrap_or(ResponseFormat::Passthrough))). Update
the call site in streaming.rs to use try_tool_response_format (or the
single-lookup inline) instead of has_exposed_tool + tool_response_format to
eliminate the duplicate lookup while preserving the same fallback behavior.
In `@protocols/src/responses.rs`:
- Around line 46-50: The struct FunctionTool currently uses
#[serde(deny_unknown_fields)] together with #[serde(flatten)] on its function
field (FunctionTool and Function), which Serde doesn't support; fix by either
removing #[serde(deny_unknown_fields)] from FunctionTool (so flatten remains) or
refactor to stop flattening: make FunctionTool contain an explicit named field
pub function: Function (remove #[serde(flatten)]), and then keep
#[serde(deny_unknown_fields)] on the inner Function type to enforce
unknown-field rejection for function tools; update the FunctionTool declaration
accordingly and ensure serde attributes are only applied in a supported
combination.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (15)
mcp/src/responses_bridge.rsmodel_gateway/src/routers/grpc/common/responses/utils.rsmodel_gateway/src/routers/grpc/harmony/builder.rsmodel_gateway/src/routers/grpc/harmony/responses/common.rsmodel_gateway/src/routers/grpc/harmony/responses/non_streaming.rsmodel_gateway/src/routers/grpc/harmony/responses/streaming.rsmodel_gateway/src/routers/grpc/harmony/stages/preparation.rsmodel_gateway/src/routers/grpc/harmony/streaming.rsmodel_gateway/src/routers/grpc/regular/responses/conversions.rsmodel_gateway/src/routers/mcp_utils.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/routers/openai/responses/utils.rsmodel_gateway/tests/api/responses_api_test.rsmodel_gateway/tests/spec/responses.rsprotocols/src/responses.rs
Description
Problem
ResponseTool was previously modeled as a flattened struct with a type field plus many optional fields shared across tool kinds. This made it easy for invalid combinations (e.g. MCP-only fields on non-MCP tools) to silently pass parsing/validation, and it doesn’t scale well as OpenAI’s Responses API has (we are going to support) more tool types and tool-specific fields. Extending support would require more and more ad-hoc, per-field validation rules to prevent “field bleed”.
Solution
Refactor ResponseTool to match the OpenAI Responses API tool schema: an internally tagged enum (type-discriminated) where each tool type owns only its valid fields. This shifts correctness from “runtime validation of a permissive struct” to “type-level correctness”, reducing validation complexity and making it straightforward to add additional OpenAI tool types in the future without proliferating conditional checks.
Changes
Test Plan
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Release Notes