fix(mcp): hide internal tool events in response streams#1405
Conversation
Signed-off-by: Daisy Zhou <zhoug9127@gmail.com>
|
Warning Rate limit exceeded
To keep reviews running without waiting, you can enable usage-based add-on for your organization. This allows additional reviews beyond the hourly cap. Account admins can enable it under billing. ⌛ 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 (1)
📝 WalkthroughWalkthroughUpdates hide internal, non-builtin MCP tool artifacts from client-facing outputs across streaming and non-streaming paths: suppressing SSE emissions during execution, filtering/redacting streaming events and response envelopes, and excluding internal servers from streamed tool listings. Documentation and tests were updated accordingly. Changes
Sequence Diagram(s)sequenceDiagram
actor Client
participant StreamingHandler
participant ToolExecutor
participant ResponseFilter
participant Storage as ToolLoopState
Client->>StreamingHandler: Submit request (may reference MCP tools)
StreamingHandler->>ToolExecutor: Start tool execution (streaming)
ToolExecutor->>ToolExecutor: Resolve tool name / is_internal_non_builtin?
alt internal non-builtin
ToolExecutor->>ToolExecutor: emit_tool_events = false
ToolExecutor->>Storage: Record call item (no SSE emitted)
else public/builtin
ToolExecutor->>StreamingHandler: Emit SSE events for client
ToolExecutor->>Storage: Record call item
end
ToolExecutor->>ResponseFilter: Execution result / events
ResponseFilter->>ResponseFilter: Check session visibility / strip_internal_mcp_artifacts
alt contains internal artifacts
ResponseFilter->>ResponseFilter: Remove tools/tool_choice/hidden items
ResponseFilter->>Client: Send redacted events/envelope
else
ResponseFilter->>Client: Send events/envelope as-is
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
There was a problem hiding this comment.
Clean, well-structured PR. The internal-tool suppression logic is correct across all three paths (tool-loop events, live streaming events, final response envelopes), and the test coverage is solid. One minor nit posted about a duplicate collect_user_function_names allocation in the streaming hot path.
Summary: 0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
There was a problem hiding this comment.
Code Review
This pull request extends the internal: true flag for MCP servers to hide internal tool details from streaming outputs, live events, and response envelopes. It introduces logic to suppress streaming tool events and redact internal artifacts from response data. The review feedback highlights a compilation error in the tests due to a missing helper function in tool_loop.rs and suggests an optimization to avoid redundant function calls in should_suppress_internal_streaming_event.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b81b530857
ℹ️ 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".
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/openai/mcp/tool_loop.rs`:
- Around line 203-205: The code currently sets emit_tool_events =
!session.is_internal_non_builtin_tool(&call.name) and then skips all
tx.send(...) paths for internal tools, which prevents early detection of
disconnected clients; change the logic so that even when emit_tool_events is
false you still perform a lightweight disconnect check before long-running MCP
calls—either by doing a non-blocking tx.try_send() of a minimal heartbeat/event
or by checking the transmit channel state (e.g., tx.is_closed() or equivalent)
and returning early on disconnect; update the same pattern around the other send
sites referenced (the blocks covering the 225-233, 248-250, and 305-313 ranges)
so every branch checks/sends a minimal probe or verifies channel liveness prior
to executing the tool call.
In `@model_gateway/src/routers/openai/responses/streaming.rs`:
- Around line 438-461: In should_suppress_internal_streaming_event, avoid
calling collect_user_function_names twice: call
collect_user_function_names(ctx.original_request) once at the top of the
function (e.g., let user_function_names =
collect_user_function_names(ctx.original_request);) and reuse that variable for
both the session.should_hide_output_item_json(item, &user_function_names) check
and the final session.is_internal_non_builtin_tool(tool_name.as_ref()) &&
!user_function_names.contains(tool_name.as_ref()) check; keep the
streaming_event_tool_name(parsed_data, handler) and session checks unchanged but
remove the second collect_user_function_names call.
🪄 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: 8642f3a2-0253-42ad-8307-ba885a392626
📒 Files selected for processing (4)
docs/reference/mcp-internal-servers.mdmodel_gateway/src/routers/openai/mcp/tool_loop.rsmodel_gateway/src/routers/openai/responses/streaming.rsmodel_gateway/src/routers/openai/responses/utils.rs
Signed-off-by: Daisy Zhou <zhoug9127@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1cc8d2315
ℹ️ 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".
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
model_gateway/src/routers/openai/responses/streaming.rs (1)
369-460: 🧹 Nitpick | 🔵 TrivialAvoid rebuilding user-function names on every streamed chunk.
should_suppress_internal_streaming_eventruns on the hot path; hoistingcollect_user_function_names(ctx.original_request)out of the helper would avoid repeatedHashSetallocations.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@model_gateway/src/routers/openai/responses/streaming.rs` around lines 369 - 460, The helper should_suppress_internal_streaming_event rebuilds the user-function HashSet each call by calling collect_user_function_names(ctx.original_request); move that call out of the hot-path and pass the precomputed set into the helper (or store it once on StreamingEventContext) so you avoid repeated HashSet allocations. Update callers of should_suppress_internal_streaming_event (the streaming loop that currently invokes it) to compute user_function_names = collect_user_function_names(ctx.original_request) once and change the helper signature to accept &HashSet<String> (or appropriate type) instead of capturing it internally; keep the existing checks (item hiding and streaming_event_tool_name) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@model_gateway/src/routers/openai/responses/streaming.rs`:
- Around line 369-460: The helper should_suppress_internal_streaming_event
rebuilds the user-function HashSet each call by calling
collect_user_function_names(ctx.original_request); move that call out of the
hot-path and pass the precomputed set into the helper (or store it once on
StreamingEventContext) so you avoid repeated HashSet allocations. Update callers
of should_suppress_internal_streaming_event (the streaming loop that currently
invokes it) to compute user_function_names =
collect_user_function_names(ctx.original_request) once and change the helper
signature to accept &HashSet<String> (or appropriate type) instead of capturing
it internally; keep the existing checks (item hiding and
streaming_event_tool_name) unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: f8af576d-8f38-40ce-b3bb-a76e92065349
📒 Files selected for processing (3)
crates/mcp/src/core/session.rsmodel_gateway/src/routers/openai/mcp/tool_loop.rsmodel_gateway/src/routers/openai/responses/streaming.rs
Signed-off-by: Daisy Zhou <zhoug9127@gmail.com>
Signed-off-by: Daisy Zhou <zhoug9127@gmail.com>
|
This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you! |
|
Hi @zhoug9127, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
Description
Problem
Internal/self-provided MCP servers can be visible to the model during Responses tool loops, but client-facing OpenAI Responses streaming and final outputs must not expose those internal tool details.
Solution
Hide internal non-builtin MCP tool artifacts from OpenAI Responses streaming and final response paths while preserving public tools and builtin-routed MCP outputs. This is PR A1 in the upstream SMG memory/LTM series: it establishes the output-privacy base that later memory provider recall PRs can reuse.
Changes
mcp_list_toolsSSE events from client-visible streaming output.Test Plan
pre-commit run --all-filescargo +nightly fmt --all --checkcargo clippy --all-targets --all-features -- -D warningscargo test --manifest-path model_gateway/Cargo.toml internal --all-featurescargo test --manifest-path model_gateway/Cargo.toml forward_streaming_event_strips_internal_tools_from_response_envelope --all-featurescargo test --manifest-path model_gateway/Cargo.toml --all-features; local run failed only inmodel_gateway/tests/otel_tracing_test.rs::test_router_with_tracingbecause the OTLP collector received 0 spans. The same exact test fails onorigin/mainin this workspace, so this appears pre-existing/local-infra rather than introduced by this PR.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Bug Fixes
Documentation
Tests