diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py index e7ab994acad..d76abefef91 100644 --- a/litellm/proxy/middleware/billable_request_metrics_middleware.py +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -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 diff --git a/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py b/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py index a5b75281491..453c4494a26 100644 --- a/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py +++ b/tests/test_litellm/proxy/middleware/test_billable_request_metrics_middleware.py @@ -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):