mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(mcp): apply post-call rewrites without stale structured output (#41530)
* fix(mcp): apply async_post_mcp_tool_call_hook content changes to the tool result Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): drop structuredContent when a post-call hook rewrites tool content Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): satisfy type discipline and result contract Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): document internal logging patch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): avoid Final assignments inside callback loops Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): run every post-call hook and chain the rewritten content Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(mcp): credit the original fix from #33403 Co-authored-by: eric <mitrecx@163.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): cover post-call logging fallback paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): cover proxy hook logging context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): preserve native structured guardrail replacements * fix(mcp): invalidate stale structure after direct content edits * fix(mcp): reconcile direct edits after callback exceptions * fix(mcp): preserve successful in-place callback rewrites --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: eric <mitrecx@163.com> Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
1c602334ee
commit
7210403e23
9 changed files with 572 additions and 40 deletions
|
|
@ -574,11 +574,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
Useful if you want to modify the standard logging payload after the MCP tool call is made.
|
||||
|
||||
To change what the caller sends back to the MCP client, mutate ``response_obj``
|
||||
in place: every call site discards the returned object, because the
|
||||
dispatcher unwraps it to ``mcp_tool_call_response`` (a raw content list, not
|
||||
a ``CallToolResult``) which the tool-call paths cannot forward. Guardrails
|
||||
that mask or reject tool output should use ``post_mcp_call`` instead.
|
||||
Modify ``mcp_tool_call_response`` in place or return a replacement response
|
||||
object to change what the caller sends back to the MCP client. Content rewrites
|
||||
discard stale structured output and mark those results as tool errors.
|
||||
Use ``post_mcp_call`` guardrails for schema-preserving structured redaction.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -221,7 +221,7 @@ from .initialize_dynamic_callback_params import (
|
|||
from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.types import EmbeddedResource, ImageContent, TextContent
|
||||
from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
|
||||
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
|
|
@ -1634,15 +1634,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
async def async_post_mcp_tool_call_hook(
|
||||
self,
|
||||
kwargs: dict,
|
||||
response_obj: Any,
|
||||
response_obj: "CallToolResult",
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
):
|
||||
"""
|
||||
Post MCP Tool Call Hook
|
||||
|
||||
Use this to modify the MCP tool call response before it is returned to the user.
|
||||
"""
|
||||
) -> "CallToolResult":
|
||||
"""Apply ordered MCP content callbacks to the result returned to the caller."""
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
|
|
@ -1650,24 +1646,51 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
dynamic_success_callbacks=self.dynamic_success_callbacks,
|
||||
global_callbacks=litellm.success_callback,
|
||||
)
|
||||
post_mcp_tool_call_response_obj: Final[MCPPostCallResponseObject] = MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=response_obj, hidden_params=HiddenParams()
|
||||
)
|
||||
hidden_params = HiddenParams()
|
||||
for callback in callbacks:
|
||||
try:
|
||||
if isinstance(callback, CustomLogger):
|
||||
response: MCPPostCallResponseObject | None = await callback.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=post_mcp_tool_call_response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
original_content = copy.deepcopy(response_obj.content)
|
||||
original_structured_content = copy.deepcopy(response_obj.structured_content)
|
||||
callback_response = MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=copy.deepcopy(original_content), hidden_params=hidden_params
|
||||
)
|
||||
######################################################################
|
||||
# if any of the callbacks modify the response, use the modified response
|
||||
# current implementation returns the first modified response
|
||||
######################################################################
|
||||
if response is not None:
|
||||
response_obj = self._parse_post_mcp_call_hook_response(response=response)
|
||||
try:
|
||||
response = await callback.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=callback_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
hook_content = (
|
||||
self._parse_post_mcp_call_hook_response(response=response)
|
||||
if response is not None
|
||||
else callback_response.mcp_tool_call_response
|
||||
)
|
||||
if response is not None:
|
||||
hidden_params = response.hidden_params
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e
|
||||
)
|
||||
hook_content = None
|
||||
structured_replacement_matches = (
|
||||
response_obj.structured_content != original_structured_content
|
||||
and (
|
||||
hook_content is None
|
||||
or hook_content == original_content
|
||||
or response_obj.content == hook_content
|
||||
)
|
||||
)
|
||||
if hook_content is not None and hook_content != original_content:
|
||||
response_obj.content[:] = hook_content
|
||||
if (
|
||||
response_obj.content != original_content
|
||||
and response_obj.structured_content is not None
|
||||
and not structured_replacement_matches
|
||||
):
|
||||
response_obj.structured_content = None
|
||||
response_obj.is_error = True
|
||||
except Exception as e:
|
||||
verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e)
|
||||
return response_obj
|
||||
|
|
|
|||
|
|
@ -2198,7 +2198,7 @@ async def _fire_mcp_tool_call_logging(
|
|||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
logging_obj.post_call(original_response=result)
|
||||
await logging_obj.async_post_mcp_tool_call_hook(
|
||||
result = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=logging_obj.model_call_details,
|
||||
response_obj=result,
|
||||
start_time=start_time,
|
||||
|
|
|
|||
|
|
@ -872,7 +872,7 @@ class LiteLLM_Proxy_MCP_Handler:
|
|||
if litellm_logging_obj:
|
||||
try:
|
||||
litellm_logging_obj.post_call(original_response=result)
|
||||
await litellm_logging_obj.async_post_mcp_tool_call_hook(
|
||||
result = await litellm_logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=litellm_logging_obj.model_call_details,
|
||||
response_obj=result,
|
||||
start_time=start_time,
|
||||
|
|
|
|||
|
|
@ -2962,7 +2962,7 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
|
|||
mcp_info={"server_name": "test_server"},
|
||||
)
|
||||
|
||||
expected_response = [TextContent(type="text", text="ok")]
|
||||
expected_response = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
|
|
@ -3038,7 +3038,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
|
|||
mcp_info={"server_name": "test_server"},
|
||||
)
|
||||
|
||||
expected_response = [TextContent(type="text", text="ok")]
|
||||
expected_response = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
|
|
|
|||
|
|
@ -5,16 +5,15 @@ import json
|
|||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from mcp.types import AudioContent, CallToolResult, ImageContent, TextContent
|
||||
from openai._legacy_response import HttpxBinaryResponseContent
|
||||
|
||||
import litellm
|
||||
|
|
@ -51,6 +50,272 @@ def logging_obj():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_mcp_tool_call_hook_preserves_and_returns_content(logging_obj):
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
class RedactingLogger(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(
|
||||
self,
|
||||
kwargs: dict[str, object],
|
||||
response_obj: MCPPostCallResponseObject,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
) -> MCPPostCallResponseObject:
|
||||
assert isinstance(response_obj.mcp_tool_call_response, list)
|
||||
assert isinstance(response_obj.mcp_tool_call_response[0], TextContent)
|
||||
response_obj.mcp_tool_call_response = [TextContent(type="text", text="[REDACTED]")]
|
||||
return response_obj
|
||||
|
||||
logging_obj.dynamic_success_callbacks = [RedactingLogger()]
|
||||
result = CallToolResult(content=[TextContent(type="text", text="SECRET-1234")], isError=False)
|
||||
|
||||
hooked_content = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=logging_obj.model_call_details,
|
||||
response_obj=result,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
)
|
||||
|
||||
assert hooked_content.content == [TextContent(type="text", text="[REDACTED]")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_mcp_tool_call_hook_chains_every_callback(logging_obj):
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
class ReplacingLogger(CustomLogger):
|
||||
def __init__(self, old: str, new: str) -> None:
|
||||
super().__init__()
|
||||
self.old: Final = old
|
||||
self.new: Final = new
|
||||
self.seen: list[str] = [] # mutable-ok: test records what each callback observed
|
||||
|
||||
async def async_post_mcp_tool_call_hook(
|
||||
self,
|
||||
kwargs: dict[str, object],
|
||||
response_obj: MCPPostCallResponseObject,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
) -> MCPPostCallResponseObject:
|
||||
first = response_obj.mcp_tool_call_response[0]
|
||||
assert isinstance(first, TextContent)
|
||||
self.seen.append(first.text)
|
||||
return MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=[TextContent(type="text", text=first.text.replace(self.old, self.new))],
|
||||
hidden_params=response_obj.hidden_params,
|
||||
)
|
||||
|
||||
first_logger: Final = ReplacingLogger("SECRET", "[S]")
|
||||
second_logger: Final = ReplacingLogger("1234", "[N]")
|
||||
logging_obj.dynamic_success_callbacks = [first_logger, second_logger]
|
||||
result = CallToolResult(content=[TextContent(type="text", text="SECRET-1234")], isError=False)
|
||||
|
||||
hooked_content = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=logging_obj.model_call_details,
|
||||
response_obj=result,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
)
|
||||
|
||||
assert first_logger.seen == ["SECRET-1234"]
|
||||
assert second_logger.seen == ["[S]-1234"]
|
||||
assert hooked_content.content == [TextContent(type="text", text="[S]-[N]")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["replace", "inplace", "empty", "inplace_none", "replace_none"])
|
||||
@pytest.mark.parametrize("structured", [False, True])
|
||||
async def test_mcp_content_rewrite_never_returns_stale_structured_data(logging_obj, mode, structured):
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
class Redactor(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time):
|
||||
block = response_obj.mcp_tool_call_response[0]
|
||||
assert isinstance(block, TextContent)
|
||||
if mode in ("inplace", "inplace_none"):
|
||||
block.text = "[REDACTED]"
|
||||
return None if mode == "inplace_none" else response_obj
|
||||
if mode == "replace_none":
|
||||
response_obj.mcp_tool_call_response = [TextContent(type="text", text="[REDACTED]")]
|
||||
return None
|
||||
return MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=[] if mode == "empty" else [TextContent(type="text", text="[REDACTED]")],
|
||||
hidden_params=HiddenParams(response_cost=0.25),
|
||||
)
|
||||
|
||||
logging_obj.dynamic_success_callbacks = [Redactor()]
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="SECRET-1234")],
|
||||
structured_content={"nested": {"secret": "SECRET-1234"}} if structured else None,
|
||||
meta={"request": "trace-1"},
|
||||
)
|
||||
logging_obj.model_call_details["original_response"] = result
|
||||
returned = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=logging_obj.model_call_details,
|
||||
response_obj=result,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
)
|
||||
assert "SECRET-1234" not in result.model_dump_json(by_alias=True)
|
||||
assert returned is result
|
||||
assert result.content == ([] if mode == "empty" else [TextContent(type="text", text="[REDACTED]")])
|
||||
assert result.structured_content is None
|
||||
assert result.is_error is structured
|
||||
assert result.meta == {"request": "trace-1"}
|
||||
assert logging_obj.model_call_details["original_response"] is result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["none", "cost", "direct", "block", "exception"])
|
||||
async def test_mcp_callbacks_preserve_effective_result_and_cost(logging_obj, mode):
|
||||
from litellm.types.llms.base import HiddenParams
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
class Callback(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time):
|
||||
if mode == "exception":
|
||||
response_obj.mcp_tool_call_response[0].text = "discarded"
|
||||
raise ValueError("non-blocking callback")
|
||||
if mode in ("direct", "block"):
|
||||
original = kwargs["original_response"]
|
||||
original.content = [TextContent(type="text", text="safe")]
|
||||
original.structured_content = {"result": "safe"}
|
||||
original.is_error = mode == "block"
|
||||
if mode == "none" or mode == "direct":
|
||||
return None
|
||||
return MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=response_obj.mcp_tool_call_response,
|
||||
hidden_params=HiddenParams(response_cost=0.25),
|
||||
)
|
||||
|
||||
logging_obj.dynamic_success_callbacks = [Callback()]
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="original")], structured_content={"result": "original"}
|
||||
)
|
||||
logging_obj.model_call_details["original_response"] = result
|
||||
returned = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=logging_obj.model_call_details,
|
||||
response_obj=result,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
)
|
||||
expected = "safe" if mode in ("direct", "block") else "original"
|
||||
assert result.content == [TextContent(type="text", text=expected)]
|
||||
assert result.structured_content == {"result": expected}
|
||||
assert result.is_error is (mode == "block")
|
||||
assert returned is result
|
||||
assert logging_obj.model_call_details.get("response_cost") == (0.25 if mode in ("cost", "block") else None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_callback_cancellation_propagates_without_mutating_result(logging_obj):
|
||||
class CancelledCallback(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time):
|
||||
response_obj.mcp_tool_call_response[0].text = "partial"
|
||||
raise asyncio.CancelledError
|
||||
|
||||
logging_obj.dynamic_success_callbacks = [CancelledCallback()]
|
||||
result = CallToolResult(content=[TextContent(type="text", text="original")])
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs={},
|
||||
response_obj=result,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
)
|
||||
assert result.content == [TextContent(type="text", text="original")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("callbacks", [[], ["prometheus"]])
|
||||
@pytest.mark.parametrize("is_error", [False, True])
|
||||
async def test_mcp_without_custom_callbacks_preserves_mixed_content(logging_obj, callbacks, is_error):
|
||||
logging_obj.dynamic_success_callbacks = callbacks
|
||||
result = CallToolResult(
|
||||
content=[
|
||||
TextContent(type="text", text="ok"),
|
||||
ImageContent(type="image", data="aW1n", mime_type="image/png"),
|
||||
AudioContent(type="audio", data="c291bmQ=", mime_type="audio/wav"),
|
||||
],
|
||||
structured_content={"result": "ok"},
|
||||
is_error=is_error,
|
||||
)
|
||||
before = result.model_dump()
|
||||
returned = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs={},
|
||||
response_obj=result,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
)
|
||||
assert returned is result
|
||||
assert returned.model_dump() == before
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("replace_structured", [False, True])
|
||||
@pytest.mark.parametrize("same_content", [False, True])
|
||||
async def test_mcp_native_structured_replacement_must_match_returned_content(
|
||||
logging_obj, replace_structured, same_content
|
||||
):
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
class NativeReplacement(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time):
|
||||
original = kwargs["original_response"]
|
||||
original.content[0].text = "native-safe"
|
||||
if replace_structured:
|
||||
original.structured_content["result"] = "native-safe"
|
||||
return MCPPostCallResponseObject(
|
||||
mcp_tool_call_response=[TextContent(type="text", text="native-safe" if same_content else "final-safe")],
|
||||
hidden_params=response_obj.hidden_params,
|
||||
)
|
||||
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="SECRET-1234")],
|
||||
structured_content={"result": "SECRET-1234"},
|
||||
)
|
||||
logging_obj.dynamic_success_callbacks = [NativeReplacement()]
|
||||
returned = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs={"original_response": result}, response_obj=result,
|
||||
start_time=datetime.datetime.now(), end_time=datetime.datetime.now(),
|
||||
)
|
||||
assert returned is result
|
||||
assert result.content == [TextContent(type="text", text="native-safe" if same_content else "final-safe")]
|
||||
assert result.structured_content == ({"result": "native-safe"} if replace_structured and same_content else None)
|
||||
assert result.is_error is not (replace_structured and same_content)
|
||||
assert "SECRET-1234" not in result.model_dump_json()
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ["none", "wrapper", "exception"])
|
||||
@pytest.mark.parametrize("structured", [False, True])
|
||||
async def test_mcp_direct_content_edit_invalidates_stale_structured_data(logging_obj, mode, structured):
|
||||
class DirectRedactor(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, end_time):
|
||||
kwargs["original_response"].content[0].text = "[REDACTED]"
|
||||
if mode == "exception":
|
||||
raise ValueError("non-blocking callback after direct edit")
|
||||
return response_obj if mode == "wrapper" else None
|
||||
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="SECRET-1234")],
|
||||
structured_content={"result": "SECRET-1234"} if structured else None,
|
||||
)
|
||||
logging_obj.dynamic_success_callbacks = [DirectRedactor()]
|
||||
returned = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs={"original_response": result}, response_obj=result,
|
||||
start_time=datetime.datetime.now(), end_time=datetime.datetime.now(),
|
||||
)
|
||||
assert returned is result
|
||||
assert result.content == [TextContent(type="text", text="[REDACTED]")]
|
||||
assert result.structured_content is None
|
||||
assert result.is_error is structured
|
||||
assert "SECRET-1234" not in result.model_dump_json()
|
||||
|
||||
|
||||
def test_get_combined_callback_list_preserves_insertion_order(logging_obj):
|
||||
assert logging_obj.get_combined_callback_list(
|
||||
dynamic_success_callbacks=["prometheus", "langfuse", "datadog", "otel", "s3"],
|
||||
|
|
|
|||
|
|
@ -8755,7 +8755,7 @@ def _call_tool_result(is_error: bool, text: str) -> CallToolResult:
|
|||
def _mock_mcp_logging_obj() -> MagicMock:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.async_post_mcp_tool_call_hook = AsyncMock()
|
||||
logging_obj.async_post_mcp_tool_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response_obj"])
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.async_failure_handler = AsyncMock()
|
||||
return logging_obj
|
||||
|
|
@ -8857,6 +8857,64 @@ async def test_fire_mcp_tool_call_logging_success_path_unchanged():
|
|||
proxy_logging_mock.post_call_failure_hook.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fire_mcp_tool_call_logging_applies_hook_content():
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_fire_mcp_tool_call_logging,
|
||||
)
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
|
||||
class RedactingLogger(CustomLogger):
|
||||
async def async_post_mcp_tool_call_hook(
|
||||
self,
|
||||
kwargs: dict[str, object],
|
||||
response_obj: MCPPostCallResponseObject,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> MCPPostCallResponseObject:
|
||||
assert isinstance(response_obj.mcp_tool_call_response, list)
|
||||
assert isinstance(response_obj.mcp_tool_call_response[0], TextContent)
|
||||
response_obj.mcp_tool_call_response = [TextContent(type="text", text="[REDACTED]")]
|
||||
return response_obj
|
||||
|
||||
logging_obj = Logging(
|
||||
model="MCP: weather/get_forecast",
|
||||
messages=[{"role": "user", "content": "tool call"}],
|
||||
stream=False,
|
||||
call_type="call_mcp_tool",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="test-mcp-hook-content",
|
||||
function_id="test-fn",
|
||||
dynamic_success_callbacks=[RedactingLogger()],
|
||||
)
|
||||
proxy_logging_mock = _mock_mcp_proxy_logging()
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="SECRET-1234")],
|
||||
structuredContent={"result": "SECRET-1234"},
|
||||
isError=False,
|
||||
)
|
||||
|
||||
with patch( # test-quality-ok: [TQ008] inject proxy logging collaborator
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj",
|
||||
proxy_logging_mock,
|
||||
):
|
||||
hooked_result = await _fire_mcp_tool_call_logging(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"),
|
||||
request_data={},
|
||||
)
|
||||
|
||||
assert isinstance(hooked_result.content[0], TextContent)
|
||||
assert hooked_result.content[0].text == "[REDACTED]"
|
||||
assert hooked_result.structured_content is None
|
||||
assert hooked_result.is_error is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fire_mcp_tool_call_logging_iserror_without_auth_skips_failure_hook():
|
||||
"""Without a UserAPIKeyAuth the failure handlers still fire but the proxy
|
||||
|
|
@ -8871,7 +8929,7 @@ async def test_fire_mcp_tool_call_logging_iserror_without_auth_skips_failure_hoo
|
|||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
|
||||
await _fire_mcp_tool_call_logging(
|
||||
logging_obj=logging_obj,
|
||||
result={"isError": True, "content": [{"type": "text", "text": "denied"}]},
|
||||
result=_call_tool_result(True, "denied"),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -827,3 +827,38 @@ class TestCiscoAIDefenseJsonRpcSuccessEnvelope:
|
|||
assert unwrapped is verdict
|
||||
else:
|
||||
assert unwrapped == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["block", "redact"])
|
||||
async def test_cisco_native_hook_through_logging_preserves_sanitized_result(action):
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
guardrail = _make_guardrail(
|
||||
inspection_type="mcp", event_hook=["pre_mcp_call", "during_mcp_call"]
|
||||
)
|
||||
verdict = (
|
||||
_violation_response(url=MCP_URL) if action == "block"
|
||||
else _redact_response(sanitized_text="[REDACTED]", url=MCP_URL)
|
||||
)
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="SECRET-1234")],
|
||||
structured_content={"result": "SECRET-1234"},
|
||||
)
|
||||
logging_obj = Logging(
|
||||
model="MCP: probe/search", messages=[], stream=False, call_type="call_mcp_tool",
|
||||
start_time=datetime.now(), litellm_call_id="cisco-hook", function_id="cisco-hook",
|
||||
dynamic_success_callbacks=[guardrail],
|
||||
)
|
||||
logging_obj.model_call_details.update({"name": "search", "arguments": {}, "original_response": result})
|
||||
with _patch_inspection_post(guardrail, AsyncMock(return_value=verdict)):
|
||||
returned = await logging_obj.async_post_mcp_tool_call_hook(
|
||||
kwargs=logging_obj.model_call_details, response_obj=result,
|
||||
start_time=datetime.now(), end_time=datetime.now(),
|
||||
)
|
||||
assert returned is result
|
||||
assert "SECRET-1234" not in returned.model_dump_json()
|
||||
assert returned.is_error is (action == "block")
|
||||
assert ("Blocked by Cisco AI Defense" if action == "block" else "[REDACTED]") in returned.content[0].text
|
||||
assert returned.structured_content == {"result": returned.content[0].text}
|
||||
|
|
|
|||
|
|
@ -1,13 +1,15 @@
|
|||
import importlib
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import types
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
from openai.types.responses.tool_param import Mcp
|
||||
import importlib
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing
|
||||
from litellm.responses import main as responses_main
|
||||
|
|
@ -15,10 +17,9 @@ from litellm.responses.mcp import litellm_proxy_mcp_handler as mcp_handler_modul
|
|||
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
from typing import Any, cast
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.types.responses.main import OutputFunctionToolCall
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
class _DummyMCPResult:
|
||||
|
|
@ -492,6 +493,157 @@ async def test_execute_tool_calls_threads_logging_obj_into_call_tool(monkeypatch
|
|||
assert call_tool_mock.await_args.kwargs["litellm_logging_obj"] is sentinel_logging_obj
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch):
|
||||
proxy_module = types.SimpleNamespace(proxy_logging_obj=None)
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module)
|
||||
|
||||
result = CallToolResult(
|
||||
content=[TextContent(type="text", text="SECRET-1234")],
|
||||
structuredContent={"result": "SECRET-1234"},
|
||||
isError=False,
|
||||
)
|
||||
fake_manager = types.SimpleNamespace(
|
||||
get_registry=MagicMock(return_value={}),
|
||||
call_tool=AsyncMock(return_value=result),
|
||||
_get_mcp_server_from_tool_name=MagicMock(return_value=None),
|
||||
get_mcp_server_by_name=MagicMock(return_value=None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
fake_manager,
|
||||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True))
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
|
||||
monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None))
|
||||
|
||||
tool_name = "deepwiki-read_wiki_structure"
|
||||
results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
|
||||
tool_server_map={tool_name: "deepwiki"},
|
||||
tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}],
|
||||
user_api_key_auth=None,
|
||||
)
|
||||
|
||||
assert results == [{"tool_call_id": "call-1", "result": "[REDACTED]", "name": tool_name}]
|
||||
assert logging_obj.async_success_handler.await_args.kwargs["result"].content[0].text == "[REDACTED]"
|
||||
assert logging_obj.async_success_handler.await_args.kwargs["result"].structured_content is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_tool_calls_returns_proxy_result_without_logging(monkeypatch):
|
||||
result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response)
|
||||
monkeypatch.setitem(
|
||||
sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj)
|
||||
)
|
||||
|
||||
fake_manager = types.SimpleNamespace(
|
||||
get_registry=MagicMock(return_value={}),
|
||||
call_tool=AsyncMock(return_value=result),
|
||||
_get_mcp_server_from_tool_name=MagicMock(return_value=None),
|
||||
get_mcp_server_by_name=MagicMock(return_value=None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
fake_manager,
|
||||
)
|
||||
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
|
||||
monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (None, None))
|
||||
|
||||
tool_name = "deepwiki-read_wiki_structure"
|
||||
results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
|
||||
tool_server_map={tool_name: "deepwiki"},
|
||||
tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}],
|
||||
user_api_key_auth=None,
|
||||
)
|
||||
|
||||
assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}]
|
||||
proxy_logging_obj.post_mcp_call_hook.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_tool_calls_passes_logging_details_to_proxy_hook(monkeypatch):
|
||||
result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response)
|
||||
monkeypatch.setitem(
|
||||
sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(proxy_logging_obj=proxy_logging_obj)
|
||||
)
|
||||
|
||||
fake_manager = types.SimpleNamespace(
|
||||
get_registry=MagicMock(return_value={}),
|
||||
call_tool=AsyncMock(return_value=result),
|
||||
_get_mcp_server_from_tool_name=MagicMock(return_value=None),
|
||||
get_mcp_server_by_name=MagicMock(return_value=None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
fake_manager,
|
||||
)
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {"request_id": "request-1"}
|
||||
logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=result)
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
|
||||
monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None))
|
||||
|
||||
tool_name = "deepwiki-read_wiki_structure"
|
||||
results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
|
||||
tool_server_map={tool_name: "deepwiki"},
|
||||
tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}],
|
||||
user_api_key_auth=None,
|
||||
)
|
||||
|
||||
assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}]
|
||||
assert proxy_logging_obj.post_mcp_call_hook.await_args.kwargs["request_data"] == logging_obj.model_call_details
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure_stage", ["post_call_hook", "success_handler"])
|
||||
async def test_execute_tool_calls_continues_when_post_call_logging_fails(monkeypatch, failure_stage: str):
|
||||
proxy_module = types.SimpleNamespace(proxy_logging_obj=None)
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_module)
|
||||
|
||||
result = CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
|
||||
fake_manager = types.SimpleNamespace(
|
||||
get_registry=MagicMock(return_value={}),
|
||||
call_tool=AsyncMock(return_value=result),
|
||||
_get_mcp_server_from_tool_name=MagicMock(return_value=None),
|
||||
get_mcp_server_by_name=MagicMock(return_value=None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
|
||||
fake_manager,
|
||||
)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
logging_obj.post_call = MagicMock()
|
||||
logging_obj.async_post_mcp_tool_call_hook = AsyncMock(
|
||||
side_effect=RuntimeError("hook failed") if failure_stage == "post_call_hook" else None,
|
||||
return_value=result,
|
||||
)
|
||||
logging_obj.async_success_handler = AsyncMock(
|
||||
side_effect=RuntimeError("success logging failed") if failure_stage == "success_handler" else None
|
||||
)
|
||||
handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler")
|
||||
monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None))
|
||||
|
||||
tool_name = "deepwiki-read_wiki_structure"
|
||||
results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
|
||||
tool_server_map={tool_name: "deepwiki"},
|
||||
tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}],
|
||||
user_api_key_auth=None,
|
||||
)
|
||||
|
||||
assert results == [{"tool_call_id": "call-1", "result": "ok", "name": tool_name}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue