feat(mcp): emit usage events for MCP tool calls - #670
Conversation
Each MCP `tools/call` now emits a UsageEvent into the same sink/pipeline as LLM usage, so MCP activity shows up in the unified usage record. The event carries who called which tool (api_key_id + new `mcp_server_name` / `mcp_tool_name`, parsed from the namespaced request), the outcome (status_code), the latency, and `inbound_protocol = "mcp"`. Rate-limited calls are recorded too (the 429 path emits before returning). The `initialize` / `tools/list` handshake and discovery methods are not metered. MCP tool calls carry no token cost yet (per-tool cost is deferred), so the token/cost fields stay zero — consistent with how other gateways log MCP calls without a per-call cost. `UsageEvent` gains two optional, skip-when-empty fields (`mcp_server_name`, `mcp_tool_name`); non-MCP events leave them empty. This is a DP-first rollout like `user_name`: older cp-api images ignore the unknown fields until the CP side adds the nullable columns (tracked follow-up). Tested: a tools/call emits one event with the right MCP attribution and zero tokens; initialize emits none. Refs AISIX-Cloud#894
|
Warning Review limit reached
Next review available in: 50 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds ChangesMCP Tool-Call Usage Metering
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/aisix-proxy/src/mcp.rs`:
- Around line 67-73: The parsed MCP server/tool names are taken directly from
request input and emitted in telemetry, so sanitize them before assignment.
Update the MCP parsing path in mcp.rs where mcp_server and mcp_tool are derived
from peek.params.name, and ensure the values are cleaned/truncated the same way
tag fields are handled in chat.rs before they are written to mcp_server_name and
mcp_tool_name, including the quota-rejection emission path referenced by the
comment.
- Around line 142-144: The MCP usage event is minting a fresh request_id inside
emit_tool_call_usage instead of using the incoming MCP correlation id, which
breaks log/row joins. Thread the handler-level request id through the /mcp flow
into emit_tool_call_usage, and use that value when building UsageEvent rather
than uuid::Uuid::new_v4(). Also make sure the MCP response echoes the same
request id so callers and proxy logs can correlate the request consistently.
🪄 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: CHILL
Plan: Pro
Run ID: b5552448-527c-4217-ae40-c2620d1f0fa3
📒 Files selected for processing (3)
crates/aisix-obs/src/usage.rscrates/aisix-proxy/src/chat.rscrates/aisix-proxy/src/mcp.rs
| let (mcp_server, mcp_tool) = if is_tool_call { | ||
| peek.as_ref() | ||
| .and_then(|p| p.params.as_ref()) | ||
| .and_then(|p| p.name.as_deref()) | ||
| .and_then(|name| name.split_once(aisix_mcp::TOOL_NAMESPACE_SEPARATOR)) | ||
| .map(|(server, tool)| (server.to_string(), tool.to_string())) | ||
| .unwrap_or_default() |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Sanitize the parsed MCP names before emitting them.
params.name is request-body input, and these values are emitted even on quota-rejection paths. Copying them verbatim into mcp_server_name / mcp_tool_name lets an authenticated caller push control characters or oversized strings into CP telemetry and OTLP export, unlike the existing tag fields in chat.rs that are sanitized before emission.
Also applies to: 149-150
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/aisix-proxy/src/mcp.rs` around lines 67 - 73, The parsed MCP
server/tool names are taken directly from request input and emitted in
telemetry, so sanitize them before assignment. Update the MCP parsing path in
mcp.rs where mcp_server and mcp_tool are derived from peek.params.name, and
ensure the values are cleaned/truncated the same way tag fields are handled in
chat.rs before they are written to mcp_server_name and mcp_tool_name, including
the quota-rejection emission path referenced by the comment.
| let event = UsageEvent { | ||
| request_id: uuid::Uuid::new_v4().to_string(), | ||
| occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== Search for request-id / x-aisix-call-id plumbing =='
fd -e rs crates | xargs rg -n -C2 'x-aisix-call-id|call_id|request_id'
echo
echo '== Inspect MCP handler context =='
sed -n '1,220p' crates/aisix-proxy/src/mcp.rsRepository: api7/aisix
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== mcp.rs outline =='
wc -l crates/aisix-proxy/src/mcp.rs
ast-grep outline crates/aisix-proxy/src/mcp.rs --view expanded
echo
echo '== mcp.rs request-id related slices =='
sed -n '1,260p' crates/aisix-proxy/src/mcp.rs
echo
echo '== request_id helper =='
sed -n '1,200p' crates/aisix-proxy/src/request_id.rs
echo
echo '== route wiring for mcp =='
rg -n -C3 'mcp\(|/mcp|mcp.rs|new_request_id\(' crates/aisix-proxy/src/lib.rs crates/aisix-proxy/src -g '!crates/aisix-proxy/src/messages.rs' -g '!crates/aisix-proxy/src/chat.rs'Repository: api7/aisix
Length of output: 25612
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo '== UsageEvent definition and request_id docs =='
rg -n -C3 'struct UsageEvent|request_id' crates -g '*.rs' | head -n 220
echo
echo '== MCP-specific response/header plumbing =='
rg -n -C3 'x-aisix-call-id|request_id|request-id|call id|call-id' crates/aisix-proxy/src/mcp.rs crates/aisix-proxy/src -g '!crates/aisix-proxy/src/messages.rs' -g '!crates/aisix-proxy/src/chat.rs' -g '!crates/aisix-proxy/src/responses.rs'
echo
echo '== Any middleware or logging that consumes x-aisix-call-id =='
rg -n -C3 'x-aisix-call-id|aisix\.request_id|request_id' crates -g '*.rs' | head -n 260Repository: api7/aisix
Length of output: 50367
Propagate the MCP request id instead of minting a new UUID.
emit_tool_call_usage currently generates its own request_id, and /mcp doesn’t thread or echo a handler-level correlation id, so these usage rows can’t be joined to the request logs/headers the way other proxy routes can. If correlation matters here, pass the request id into this helper and include it on the response.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/aisix-proxy/src/mcp.rs` around lines 142 - 144, The MCP usage event is
minting a fresh request_id inside emit_tool_call_usage instead of using the
incoming MCP correlation id, which breaks log/row joins. Thread the
handler-level request id through the /mcp flow into emit_tool_call_usage, and
use that value when building UsageEvent rather than uuid::Uuid::new_v4(). Also
make sure the MCP response echoes the same request id so callers and proxy logs
can correlate the request consistently.
Audit (CLAUDE.md §8) on #670 returned a MEDIUM: the "rate-limited calls are recorded too" claim had dedicated emit-before-return code but no test (the existing 429 test ran with a disabled sink). Added `rate_limited_tool_call_still_emits_usage_event` (rpm=1: the 429'd call emits a UsageEvent with status_code=429 + mcp attribution), and an exactly-one-event assertion to the happy-path test. Refs AISIX-Cloud#894
Independent audit (CLAUDE.md §8): BLOCK → resolvedAuditor built/tested/clippy'd in an isolated worktree — all green (509 tests). Verified: exactly one emit per tool call (429 path emits then returns → post-gateway emit unreachable; success path emits once; initialize/tools-list emit zero); One MEDIUM — fixed in the pushed commit:
Non-blocking:
|
1 participant
What
DP-4 slice ⑤b (observability). Each MCP
tools/callnow emits a UsageEvent into the same sink/pipeline as LLM usage, so MCP activity lands in the unified usage record — matching how the leading gateways log MCP calls (Kong: MCP audit log + metrics; Portkey: unified MCP+LLM traces), framed AISIX's way: the same usage pipeline, not a separate MCP surface.How
/mcphandler peeks the JSON-RPCparams.name(it already buffers the body for rate-limit), splits the namespaced<server>__<tool>, and after the gateway call emits aUsageEventviastate.usage_sink.try_emit("mcp", ..)(non-blocking, fire-and-forget — same as LLM handlers).api_key_id,status_code,latency_ms,inbound_protocol="mcp", and two newmcp_server_name/mcp_tool_name. Rate-limited calls are recorded too (the 429 path emits before returning).initialize/tools/listare not metered.UsageEventgains two optional, skip-when-empty fields. Non-MCP events leave them empty (the chat path now uses..Default::default()). DP-first rollout likeuser_name: older cp-api ignores the unknown fields until the CP migration lands — filed as AISIX-Cloud#908.Reference implementations (CLAUDE.md §7)
Re-verified against primary sources (2026-06-30): all three mainstream gateways log MCP calls — LiteLLM (same spend-log table,
mcp_namespaced_tool_name,call_type=call_mcp_tool), Kong (MCP audit log + Prometheus metrics, 3.12+), Portkey (unified MCP+LLM logs/traces). So MCP usage logging is table-stakes; AISIX's framing is unifying it into the sameUsageEventas LLM. Per-tool cost: only LiteLLM has it — we defer, like Kong/Portkey.Test plan
emits_usage_event_for_tool_call_only: atools/callemits exactly one event withinbound_protocol="mcp",mcp_server_name/mcp_tool_nameparsed from the namespaced name, the rightapi_key_id, and zero tokens;initializeemits none. (Injects a realUsageSinkwith a receiver and asserts on the channel.)fmt,clippy --workspace --all-targets -D warnings,cargo test -p aisix-obs -p aisix-proxygreen.Remaining
Slice ⑤a — guardrails over MCP (the differentiator: Kong won't, Portkey "coming soon", LiteLLM experimental+buggy) — is the next and final DP-4 PR.
Refs AISIX-Cloud#894
Summary by CodeRabbit
New Features
Bug Fixes