Conversation
… policy Signed-off-by: Daisy Zhou <zhoug9127@gmail.com>
…ity logic Signed-off-by: Daisy Zhou <zhoug9127@gmail.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ 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 |
|
Superseded by #1361 (same changes opened from upstream branch instead of fork head). Closing this one to avoid duplicate review threads. |
github-actions
Bot
added
documentation
![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 updates the MCP tool interception and redaction logic to support streaming responses and mixed-mode tool execution (where both MCP and user-defined function tools are present). It introduces refined visibility checks for streaming output items and ensures that internal MCP tools are correctly redacted across both non-streaming and streaming paths. My feedback highlights a 'fail-open' security concern in the redaction logic, suggests an optimization for cloning in the dedupe key generation, and points out an overly restrictive filter in the streaming redaction logic that could lead to information leakage.
| output.retain(|item| { | ||
| let Ok(json) = to_value(item) else { | ||
| return true; | ||
| }; | ||
| !session.should_hide_output_item_json(&json, user_function_names) | ||
| }); |
There was a problem hiding this comment.
The redaction logic here is 'fail-open'. If to_value(item) fails, the item is retained and potentially leaked to the client. For security and privacy features like redaction, it is safer to 'fail-closed' by returning false on error. Additionally, as per repository guidelines, failures in serialization should be logged as warnings to aid in debugging rather than failing silently.
output.retain(|item| {
match to_value(item) {
Ok(json) => !session.should_hide_output_item_json(&json, user_function_names),
Err(e) => {
tracing::warn!(error = ?e, "Failed to serialize item for redaction check; hiding item");
false
}
}
});References
- Instead of silently ignoring potential failures (e.g., from serialization), log them as warnings to aid in debugging. In Rust, prefer using unwrap_or_else to log an error over unwrap_or_default which would fail silently.
| } | ||
|
|
||
| let server_label = item.get("server_label").and_then(|value| value.as_str())?; | ||
| let tools = item.get("tools").cloned().unwrap_or_else(|| json!([])); |
There was a problem hiding this comment.
Avoid unnecessary cloning of the tools array. Since mcp_list_tools_dedupe_key takes a reference to Value, you can pass the reference directly.
| let tools = item.get("tools").cloned().unwrap_or_else(|| json!([])); | |
| let tools = item.get("tools").unwrap_or(&Value::Null); |
| Some(OutputItemEvent::ADDED) | Some(OutputItemEvent::DONE) => parsed_data | ||
| .get("item") | ||
| .filter(|item| { | ||
| item.get("type") | ||
| .and_then(|v| v.as_str()) | ||
| .is_some_and(is_function_call_type) | ||
| }) | ||
| .is_some_and(|item| session.should_hide_output_item_json(item, user_function_names)), |
There was a problem hiding this comment.
The filter is_function_call_type is too restrictive. should_hide_output_item_json is designed to handle multiple item types. By restricting this check to function calls, other internal output items present in the final response (e.g., if an internal mcp_list_tools item is received from an upstream gateway) might be leaked to the client. Ensure redaction logic targets all relevant variants present in the output.
Some(OutputItemEvent::ADDED) | Some(OutputItemEvent::DONE) => parsed_data
.get("item")
.is_some_and(|item| session.should_hide_output_item_json(item, user_function_names)),References
- When redacting items from a response, ensure the redaction logic targets only the item variants that are actually present in the final response output. Intermediate item types that are not part of the final assembled response do not need redaction.
1 participant
Description
Problem
Streaming MCP interception had a mixed-batch gap: when a streamed turn contained both user function calls and MCP-intercepted calls, final
response.completed/persisted output could omit already-executed MCP output items. That creates client state mismatch and replay risk.In addition, output visibility/index behavior relied on subtle ordering assumptions with limited local guardrails/tests.
Solution
function_callplaceholders and injecting executed MCP output items beforeresponse.completedemission and persistence.resolve_output_index_for_forwardingside effects and interception vs redaction context usage.mark_output_hiddenis called after an index is already visible.output_item.donehandling,Changes
model_gateway/src/routers/openai/responses/streaming.rssend_final_response_eventand persisted final response path when interception is disabled but MCP calls executed.model_gateway/src/routers/openai/mcp/tool_handler.rsmodel_gateway/src/routers/openai/mcp/tool_loop.rsmodel_gateway/src/routers/openai/mcp/mod.rsmodel_gateway/src/routers/grpc/regular/responses/streaming.rsmodel_gateway/src/routers/grpc/harmony/streaming.rsTest Plan
Reproducible verification run:
cargo test -p smg --lib routers::openai::mcp::tool_handler::tests::cargo test -p smg --lib routers::openai::responses::streaming::tests::cargo test -p smg --lib routers::openai::mcp::tool_loop::tests::cargo test -p smg --test api_tests responsescargo test -p smg --test api_tests mcppre-commit run --all-filesExpected/observed: all commands pass.
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses