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

fix: fire spend log callbacks from MCP REST /tools/call endpoint

Browse source 
The REST endpoint (/mcp-rest/tools/call) was not firing async_success_handler
after tool execution, so MCP tool calls made via REST never created spend logs.
Added the same success handler + post_call flow that the SSE/MCP protocol path
uses (server.py call_mcp_tool), ensuring spend logs are written and the MCP
server index is populated for operational visibility.

Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
This commit is contained in:
Cursor Agent 2026-03-06 04:58:09 +00:00
parent 2dfa8c266e
commit d8d232e285
1 changed files with 23 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -506,18 +506,39 @@ if MCP_AVAILABLE:
)
# Call execute_mcp_tool directly (permission checks already done)
start_time = datetime.now()
litellm_logging_obj_for_call = data.get("litellm_logging_obj")
result = await execute_mcp_tool(
name=tool_name,
arguments=tool_arguments,
allowed_mcp_servers=allowed_mcp_servers,
start_time=datetime.now(),
start_time=start_time,
user_api_key_auth=data.get("user_api_key_auth"),
mcp_auth_header=data.get("mcp_auth_header"),
mcp_server_auth_headers=data.get("mcp_server_auth_headers"),
oauth2_headers=data.get("oauth2_headers"),
raw_headers=data.get("raw_headers"),
litellm_logging_obj=data.get("litellm_logging_obj"),
litellm_logging_obj=litellm_logging_obj_for_call,
)
if litellm_logging_obj_for_call:
litellm_logging_obj_for_call.post_call(original_response=result)
end_time = datetime.now()
await litellm_logging_obj_for_call.async_post_mcp_tool_call_hook(
kwargs=litellm_logging_obj_for_call.model_call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
)
litellm_logging_obj_for_call.call_type = (
CallTypes.call_mcp_tool.value
)
await litellm_logging_obj_for_call.async_success_handler(
result=result,
start_time=start_time,
end_time=end_time,
)
return result
except BlockedPiiEntityError as e:
verbose_logger.error(f"BlockedPiiEntityError in MCP tool call: {str(e)}")