scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 186

fix(proxy): bill mcp-rest tool calls and bare a2a agent invokes

Browse source 
POST /mcp-rest/tools/call executes a tool and fires the same MCP spend
logging as the /mcp transport, and POST /a2a/{agent_id} is the JSON-RPC
invoke route whose method (message/send or message/stream) travels in
the body; both returned 2xx without being recorded
This commit is contained in:
Yassin Kortam 2026-07-09 11:09:18 +03:00
parent b25c2d2999
commit 3271c0f7d4
2 changed files with 18 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/proxy/middleware/billable_request_metrics_middleware.py
View file

@ -97,9 +97,16 @@ def _classify_llm_route(path: str) -> Optional[str]:
_MCP_MANAGEMENT_PREFIX = "/v1/mcp"
_MCP_DYNAMIC_TRANSPORT = re.compile(r"/(?:toolset/)?[^/]+/mcp")
# The REST wrapper's tool-call endpoint executes a tool and fires the same MCP
# spend logging as the /mcp transport; its list/test siblings do not bill.
_MCP_REST_TOOL_CALL = "/mcp-rest/tools/call"
_A2A_INVOKE_SUFFIX = "/message/send"
_A2A_TRANSPORT_PREFIXES: tuple[str, ...] = ("/v1/a2a/", "/a2a/")
# Bare POST /a2a/{agent_id} is the JSON-RPC invoke route (method message/send
# or message/stream travels in the body). /v1/a2a/discover has an extra path
# segment before the agent id, so this pattern cannot match it.
_A2A_BARE_INVOKE = re.compile(r"/a2a/[^/]+")
def _classify_mcp_route(path: str) -> Optional[str]:
@ -107,6 +114,8 @@ def _classify_mcp_route(path: str) -> Optional[str]:
return None
if path == "/mcp" or path.startswith("/mcp/"):
return "/mcp"
if path == _MCP_REST_TOOL_CALL:
return "/mcp"
if _MCP_DYNAMIC_TRANSPORT.fullmatch(path) is not None:
return "/mcp"
return None
@ -115,6 +124,8 @@ def _classify_mcp_route(path: str) -> Optional[str]:
def _classify_a2a_route(path: str) -> Optional[str]:
if path.endswith(_A2A_INVOKE_SUFFIX) and any(path.startswith(prefix) for prefix in _A2A_TRANSPORT_PREFIXES):
return "/a2a"
if _A2A_BARE_INVOKE.fullmatch(path) is not None:
return "/a2a"
return None

8
tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py
View file

@ -116,8 +116,13 @@ def test_is_pure_asgi_not_base_http_middleware():
("/github/mcp/", (BillableCategory.MCP, "/mcp")),
("/toolset/my-tools/mcp", (BillableCategory.MCP, "/mcp")),
("/github,slack/mcp", (BillableCategory.MCP, "/mcp")),
# REST wrapper tool execution fires the same MCP spend logging as /mcp
("/mcp-rest/tools/call", (BillableCategory.MCP, "/mcp")),
("/a2a/agent-1/message/send", (BillableCategory.A2A, "/a2a")),
("/v1/a2a/agent-9/message/send", (BillableCategory.A2A, "/a2a")),
# bare invoke route: the JSON-RPC method (message/send or message/stream)
# travels in the body, not the path
("/a2a/agent-1", (BillableCategory.A2A, "/a2a")),
],
)
def test_classify_billable(path: str, expected: Tuple[BillableCategory, str]):
@ -160,7 +165,8 @@ def test_classify_billable(path: str, expected: Tuple[BillableCategory, str]):
"/v1/a2a/discover",
"/.well-known/oauth-protected-resource/github/mcp",
"/mcp-rest/tools/list",
"/mcp-rest/tools/call",
"/mcp-rest/test/connection",
"/mcp-rest/test/tools/list",
],
)
def test_classify_non_billable_returns_none(path: str):