mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(mcp): enforce code output policies and interaction boundaries
This commit is contained in:
moe-berri
parent
dfa4936998
commit
55eeadafa7
9 changed files with 488 additions and 20 deletions
|
|
@ -18,6 +18,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.mcp_server.contracts import OperationContext
|
||||
from litellm.proxy.mcp_server.legacy_callbacks import noninteractive_code_execution
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.mcp import MCPCodeModeSettings
|
||||
|
||||
|
|
@ -206,6 +207,7 @@ async def cleanup_cancelled_call(data: Mapping[str, object], caller: UserAPIKeyA
|
|||
async def run_code_tool(arguments: Mapping[str, object] | None, context: OperationContext) -> CallToolResult:
|
||||
from litellm.proxy.mcp_server.operations import (
|
||||
_fire_mcp_tool_call_logging,
|
||||
_run_post_mcp_call_guardrails,
|
||||
fire_mcp_tool_call_failure_logging,
|
||||
prepare_virtual_mcp_call,
|
||||
)
|
||||
|
|
@ -235,7 +237,14 @@ async def run_code_tool(arguments: Mapping[str, object] | None, context: Operati
|
|||
processed: Final = TypeAdapter(dict[str, object]).validate_python(
|
||||
data.get("modified_arguments", data.get("mcp_arguments", data.get("arguments", {})))
|
||||
)
|
||||
result: Final = await _execute_script(processed, traced)
|
||||
executed: Final = await _execute_script(processed, traced)
|
||||
result: Final = await _run_post_mcp_call_guardrails(
|
||||
executed,
|
||||
litellm_logging_obj=logging_obj,
|
||||
original_auth=caller,
|
||||
guardrail_auth=caller,
|
||||
request_data=data,
|
||||
)
|
||||
if logging_obj is None:
|
||||
return result
|
||||
return await _fire_mcp_tool_call_logging(
|
||||
|
|
@ -360,7 +369,8 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon
|
|||
)
|
||||
return serialize_tool_result(result)
|
||||
|
||||
result: Final = await execute_code(request.code, call_tool)
|
||||
with noninteractive_code_execution():
|
||||
result: Final = await execute_code(request.code, call_tool)
|
||||
return CallToolResult(
|
||||
content=[
|
||||
TextContent(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
|
@ -26,6 +25,7 @@ from litellm.proxy.mcp_server.code_worker import (
|
|||
Reply,
|
||||
Start,
|
||||
encode_message,
|
||||
json_bytes,
|
||||
load_json,
|
||||
)
|
||||
|
||||
|
|
@ -117,7 +117,7 @@ async def _exchange(
|
|||
calls.start(message.id)
|
||||
group.start_soon(dispatch, message)
|
||||
case Done():
|
||||
if len(json.dumps(message.output, ensure_ascii=False).encode("utf-8")) > settings.output_bytes:
|
||||
if len(json_bytes(message.output)) > settings.output_bytes:
|
||||
return CodeExecutionResult(
|
||||
error="Code output exceeds its byte limit", calls_started=len(calls.started)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -82,10 +82,13 @@ def load_json(value: str) -> JsonValue:
|
|||
return JSON_VALUE.validate_python(json.loads(value, parse_int=_integer, parse_float=_float, parse_constant=_float))
|
||||
|
||||
|
||||
def json_bytes(value: object) -> bytes:
|
||||
return json.dumps(value, allow_nan=False, ensure_ascii=False).encode("utf-8", errors="backslashreplace")
|
||||
|
||||
|
||||
def encode_message(value: WireModel, maximum: int) -> bytes:
|
||||
text: Final = json.dumps(value.model_dump(), allow_nan=False, ensure_ascii=False)
|
||||
load_json(text)
|
||||
encoded: Final = (text + "\n").encode("utf-8")
|
||||
encoded: Final = json_bytes(value.model_dump()) + b"\n"
|
||||
load_json(encoded.decode("utf-8"))
|
||||
if len(encoded) > maximum:
|
||||
raise ValueError("Code execution message exceeds its byte limit")
|
||||
return encoded
|
||||
|
|
@ -232,13 +235,13 @@ def _run(start: Start) -> Done:
|
|||
receive: Final = runtime.eval(_BRIDGE)
|
||||
if not isinstance(receive, _Receiver):
|
||||
raise TypeError("Invalid JavaScript bridge")
|
||||
receive("start", start.model_dump_json())
|
||||
receive("start", json_bytes(start.model_dump()).decode("utf-8"))
|
||||
while output.result is None:
|
||||
while runtime.execute_pending_job():
|
||||
if output.result is not None:
|
||||
return output.result
|
||||
if output.result is None:
|
||||
receive("reply", _read_reply(start.limits.message_bytes).model_dump_json())
|
||||
receive("reply", json_bytes(_read_reply(start.limits.message_bytes).model_dump()).decode("utf-8"))
|
||||
return output.result
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
from collections.abc import Mapping
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx2
|
||||
|
|
@ -22,6 +24,17 @@ from litellm.proxy.mcp_server.scope_challenge import (
|
|||
scope_response_from_request,
|
||||
)
|
||||
|
||||
_code_execution_active: Final[ContextVar[bool]] = ContextVar("mcp_code_execution_active", default=False)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def noninteractive_code_execution() -> Generator[None, None, None]:
|
||||
token: Final = _code_execution_active.set(True)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_code_execution_active.reset(token)
|
||||
|
||||
|
||||
class SamplingCallback(Protocol):
|
||||
async def __call__(
|
||||
|
|
@ -77,6 +90,7 @@ def create_sampling_callback(
|
|||
def create_elicitation_callback(timeout: float | None = None) -> ElicitationCallback:
|
||||
from litellm.proxy.mcp_server.server import get_active_mcp_session
|
||||
|
||||
code_execution_active: Final = _code_execution_active.get()
|
||||
downstream_session: Final = get_active_mcp_session()
|
||||
request: Final = get_active_mcp_request_ctx()
|
||||
client_params: Final = downstream_session.client_params if downstream_session is not None else None
|
||||
|
|
@ -89,6 +103,14 @@ def create_elicitation_callback(timeout: float | None = None) -> ElicitationCall
|
|||
relay_timeout: Final = timeout if timeout is not None else MCP_CLIENT_TIMEOUT
|
||||
|
||||
async def callback(context: object, params: ElicitRequestParams) -> ElicitResult | ErrorData:
|
||||
if code_execution_active:
|
||||
return ErrorData(
|
||||
code=-32602,
|
||||
message=(
|
||||
"Interactive elicitation is unsupported in code mode. Call this tool directly instead. "
|
||||
"Earlier actions may already have completed; do not automatically replay the script."
|
||||
),
|
||||
)
|
||||
if request is not None and request.protocol_version == "2026-07-28":
|
||||
return ErrorData(
|
||||
code=-32602,
|
||||
|
|
|
|||
|
|
@ -1604,3 +1604,165 @@ def test_legacy_notification_dispatch_finishes_before_session_deletion() -> None
|
|||
assert terminated.status_code == 200
|
||||
finally:
|
||||
deleted.set()
|
||||
|
||||
|
||||
def test_code_mode_blocks_legacy_elicitation_but_preserves_ordinary_consent(tmp_path: Path) -> None:
|
||||
import asyncio
|
||||
|
||||
import httpx2
|
||||
from integration._support.mcp import mcp_peer, tool_calls
|
||||
from mcp import ClientSession
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.types import ElicitResult
|
||||
|
||||
with mcp_peer("sse", rich=True) as peer, httpx.Client() as http_client:
|
||||
config = tmp_path / "proxy.yaml"
|
||||
config.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [],
|
||||
"mcp_servers": {
|
||||
"peer": {
|
||||
"url": peer.url,
|
||||
"transport": "sse",
|
||||
"protocol_version": "2025-11-25",
|
||||
"allow_elicitation": True,
|
||||
}
|
||||
},
|
||||
"general_settings": {
|
||||
"master_key": _KEY,
|
||||
"store_model_in_db": False,
|
||||
"mcp_advertised_versions": ["2025-11-25"],
|
||||
},
|
||||
"litellm_settings": {"mcp_code_mode": True},
|
||||
}
|
||||
)
|
||||
)
|
||||
consent_requests: list[str] = []
|
||||
|
||||
async def consent(context, params):
|
||||
consent_requests.append(params.message)
|
||||
return ElicitResult(action="accept", content={"confirmed": True})
|
||||
|
||||
async def exercise(gateway):
|
||||
async with (
|
||||
httpx2.AsyncClient(headers={"Authorization": "Bearer " + gateway.key}) as client,
|
||||
streamable_http_client(
|
||||
str(gateway.client.base_url).rstrip("/") + "/mcp/", http_client=client
|
||||
) as streams,
|
||||
ClientSession(streams[0], streams[1], elicitation_callback=consent) as session,
|
||||
):
|
||||
await session.initialize()
|
||||
await session.list_tools()
|
||||
ordinary = await session.call_tool("peer-elicit", {"question": "Confirm"})
|
||||
assert not ordinary.is_error and ordinary.content[0].text == "elicited:accept", ordinary
|
||||
assert consent_requests == ["Confirm"]
|
||||
consent_requests.clear()
|
||||
nested = await session.call_tool(
|
||||
"litellm_execute_code",
|
||||
{"code": 'text(await tools.call_tool({tool_id:"peer-elicit",arguments:{question:"Confirm"}}))'},
|
||||
)
|
||||
assert consent_requests == [], nested
|
||||
body = json.loads(nested.content[0].text)
|
||||
assert body["calls_started"] == 1, body
|
||||
child = body["output"][0]
|
||||
assert child["isError"] is True, child
|
||||
assert "elicitation" in json.dumps(child).lower(), child
|
||||
assert "unsupported" in json.dumps(child).lower(), child
|
||||
# The request-local restriction must not leak into a subsequent ordinary call.
|
||||
recovered = await session.call_tool("peer-elicit", {"question": "Confirm again"})
|
||||
assert not recovered.is_error and recovered.content[0].text == "elicited:accept", recovered
|
||||
assert consent_requests == ["Confirm again"]
|
||||
|
||||
with owned_proxy(
|
||||
Gateway(http_client, _KEY, peer.url),
|
||||
tmp_path / "gateway",
|
||||
{
|
||||
"PYTHONPATH": _PROXY_PYTHONPATH,
|
||||
"STORE_MODEL_IN_DB": "False",
|
||||
"DISABLE_SCHEMA_UPDATE": "true",
|
||||
"LITELLM_SALT_KEY": "code-elicitation-test",
|
||||
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
|
||||
},
|
||||
config=config,
|
||||
workers=1,
|
||||
database_setup=(),
|
||||
remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL", "REDIS_URL", "REDIS_HOST"),
|
||||
) as gateway:
|
||||
asyncio.run(exercise(gateway))
|
||||
assert len(tool_calls(peer.drain())) == 3, "Consent handling restarted an upstream tool"
|
||||
|
||||
|
||||
def test_code_mode_checks_final_output_guardrail_on_all_routes(tmp_path: Path) -> None:
|
||||
from integration._support.mcp import mcp_peer, official_client_outcomes, tool_calls
|
||||
|
||||
(tmp_path / "code_output_guard.py").write_text("""
|
||||
from fastapi import HTTPException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
||||
class RejectOutput(CustomGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
if any("qa-rejected-output" in text for text in inputs.get("texts", [])):
|
||||
raise HTTPException(status_code=400, detail="QA output rejected")
|
||||
return inputs
|
||||
""")
|
||||
with mcp_peer() as peer, httpx.Client() as client:
|
||||
config = tmp_path / "proxy.yaml"
|
||||
config.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [],
|
||||
"mcp_servers": {"peer": {"url": peer.url, "transport": "http"}},
|
||||
"general_settings": {"master_key": _KEY, "store_model_in_db": False},
|
||||
"litellm_settings": {"mcp_code_mode": True},
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": "qa-output",
|
||||
"litellm_params": {
|
||||
"guardrail": "code_output_guard.RejectOutput",
|
||||
"mode": "post_mcp_call",
|
||||
"default_on": True,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
with owned_proxy(
|
||||
Gateway(client, _KEY, peer.url),
|
||||
tmp_path / "gateway",
|
||||
{
|
||||
"PYTHONPATH": os.pathsep.join([str(tmp_path), _PROXY_PYTHONPATH]),
|
||||
"STORE_MODEL_IN_DB": "False",
|
||||
"DISABLE_SCHEMA_UPDATE": "true",
|
||||
"LITELLM_SALT_KEY": "code-output-test",
|
||||
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
|
||||
},
|
||||
config=config,
|
||||
workers=2,
|
||||
database_setup=(),
|
||||
remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL", "REDIS_URL", "REDIS_HOST"),
|
||||
) as gateway:
|
||||
outcomes = []
|
||||
for path in ("/mcp", "/mcp/proxy", "rest"):
|
||||
responses = []
|
||||
for code in ('text("qa-rejected-output")', 'text("qa-allowed-output")'):
|
||||
if path == "rest":
|
||||
result = McpCaller(gateway, gateway.key, "rest").call("litellm_execute_code", {"code": code})
|
||||
else:
|
||||
_, result = official_client_outcomes(
|
||||
gateway, gateway.key, path, "litellm_execute_code", {"code": code}
|
||||
)
|
||||
responses.append(result)
|
||||
outcomes.append((path, *responses))
|
||||
for path, rejected, allowed in outcomes:
|
||||
assert not rejected.ok, (path, rejected)
|
||||
assert "QA output rejected" in (rejected.error or ""), (path, rejected)
|
||||
assert allowed.ok and allowed.text is not None, (path, allowed)
|
||||
assert json.loads(allowed.text)["output"] == ["qa-allowed-output"]
|
||||
surrogate = McpCaller(gateway, gateway.key, "rest").call(
|
||||
"litellm_execute_code", {"code": "text('kept'); text(String.fromCharCode(0xd800));"}
|
||||
)
|
||||
assert surrogate.ok and surrogate.text is not None, surrogate
|
||||
assert json.loads(surrogate.text) == {"output": ["kept", chr(0xD800)], "error": None, "calls_started": 0}
|
||||
assert tool_calls(peer.drain()) == (), "Pure code unexpectedly invoked an upstream tool"
|
||||
|
|
|
|||
|
|
@ -35,6 +35,22 @@ async def test_code_roundtrips_full_results_and_chains_calls() -> None:
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_preserves_surrogate_output_and_tool_values() -> None:
|
||||
result: Final = await execute_code(
|
||||
"text('kept'); text(String.fromCharCode(0xd800)); "
|
||||
"text(await tools.call_tool({tool_id:'echo', arguments:{value:String.fromCharCode(0xdc00)}}));",
|
||||
_echo,
|
||||
)
|
||||
assert result.error is None
|
||||
assert result.calls_started == 1
|
||||
assert result.output == (
|
||||
"kept",
|
||||
"\ud800",
|
||||
{"structuredContent": {"value": "\udc00"}, "content": [], "isError": False},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_script_near_byte_limit_survives_json_escaping() -> None:
|
||||
source = "/*" + '\\"' * ((CodeExecutionLimits().script_bytes - 20) // 2) + "*/text(42)"
|
||||
|
|
|
|||
|
|
@ -1,18 +1,39 @@
|
|||
import asyncio
|
||||
import subprocess
|
||||
from collections.abc import Mapping, Sequence
|
||||
from contextvars import Context
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
from anyio.abc import ByteReceiveStream, ByteSendStream, Process
|
||||
from mcp.types import TextContent, Tool
|
||||
from fastapi import HTTPException
|
||||
from mcp.server.connection import Connection
|
||||
from mcp.server.context import ServerRequestContext
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.shared.dispatcher import CallOptions
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from mcp.shared.transport_context import TransportContext
|
||||
from mcp.types import (
|
||||
ClientCapabilities,
|
||||
ElicitationCapability,
|
||||
ElicitRequestFormParams,
|
||||
ElicitResult,
|
||||
ErrorData,
|
||||
FormElicitationCapability,
|
||||
Implementation,
|
||||
InitializeRequestParams,
|
||||
TextContent,
|
||||
Tool,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
|
@ -20,14 +41,18 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelReq
|
|||
from litellm.proxy.mcp_server.code_execution import CODE_TOOL_NAME, run_code_tool
|
||||
from litellm.proxy.mcp_server.code_worker import Call, Done, Reply, Start, encode_message, load_json
|
||||
from litellm.proxy.mcp_server.contracts import OperationContext
|
||||
from litellm.proxy.mcp_server.legacy_callbacks import create_elicitation_callback
|
||||
from litellm.proxy.mcp_server.mcp_context import active_mcp_request_ctx_var
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
from litellm.proxy.mcp_server.result_conversion import INPUT_REQUIRED_UNSUPPORTED_MESSAGE
|
||||
from litellm.proxy.mcp_server.server import active_mcp_session_var
|
||||
from litellm.proxy.mcp_server.tool_registry import global_mcp_tool_registry
|
||||
from litellm.proxy.mcp_server.tool_search import mcp_proxy_tool_id, with_mcp_proxy_identity
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.mcp import MCPPostCallResponseObject, MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures("mcp_code_runtime")
|
||||
|
||||
|
|
@ -35,8 +60,9 @@ pytestmark: Final = pytest.mark.usefixtures("mcp_code_runtime")
|
|||
class ScriptedProcess(Process, ByteReceiveStream, ByteSendStream):
|
||||
"""Only the external process transport is replaced; each reply comes from LiteLLM."""
|
||||
|
||||
def __init__(self, calls: tuple[Call, ...]) -> None:
|
||||
def __init__(self, calls: tuple[Call, ...], output: tuple[JsonValue, ...] | None = None) -> None:
|
||||
self.calls = calls
|
||||
self.output = output
|
||||
self.started: Start | None = None
|
||||
self.replies: tuple[Reply, ...] = ()
|
||||
self.stopped = False
|
||||
|
|
@ -53,7 +79,9 @@ class ScriptedProcess(Process, ByteReceiveStream, ByteSendStream):
|
|||
self.calls[len(self.replies)]
|
||||
if len(self.replies) < len(self.calls)
|
||||
else Done(
|
||||
output=tuple(
|
||||
output=self.output
|
||||
if self.output is not None
|
||||
else tuple(
|
||||
dict(reply.value) if reply.value is not None else {"error": reply.error} for reply in self.replies
|
||||
)
|
||||
)
|
||||
|
|
@ -103,8 +131,10 @@ class ScriptedProcess(Process, ByteReceiveStream, ByteSendStream):
|
|||
return None
|
||||
|
||||
|
||||
def install_transport(monkeypatch: pytest.MonkeyPatch, calls: tuple[Call, ...]) -> ScriptedProcess:
|
||||
process: Final = ScriptedProcess(calls)
|
||||
def install_transport(
|
||||
monkeypatch: pytest.MonkeyPatch, calls: tuple[Call, ...], output: tuple[JsonValue, ...] | None = None
|
||||
) -> ScriptedProcess:
|
||||
process: Final = ScriptedProcess(calls, output)
|
||||
|
||||
async def open_process(command: Sequence[str], *, cwd: str, env: Mapping[str, str], stderr: int) -> Process:
|
||||
assert "-I" in command and "-u" in command
|
||||
|
|
@ -416,3 +446,197 @@ async def test_successful_continuation_text_is_not_reclassified_as_an_error(
|
|||
"isError": False,
|
||||
"resultType": "complete",
|
||||
}
|
||||
|
||||
|
||||
class PromptChannel:
|
||||
can_send_request = True
|
||||
request_id = "ordinary-request"
|
||||
message_metadata = None
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.transport = TransportContext(kind="in-process", can_send_request=True)
|
||||
self.cancel_requested = anyio.Event()
|
||||
self.prompts: tuple[Mapping[str, object], ...] = ()
|
||||
|
||||
async def send_raw_request(
|
||||
self, method: str, params: Mapping[str, object] | None, opts: CallOptions | None = None
|
||||
) -> dict[str, JsonValue]: # mutable-ok: MCP SDK outbound protocol requires a dictionary
|
||||
assert method == "elicitation/create" and params is not None
|
||||
self.prompts = (*self.prompts, dict(params))
|
||||
return {"action": "accept", "content": {"value": "accepted"}}
|
||||
|
||||
async def notify(self, method: str, params: Mapping[str, object] | None, opts: CallOptions | None = None) -> None:
|
||||
pytest.fail("Elicitation must not send unrelated notifications")
|
||||
|
||||
async def progress(self, progress: float, total: float | None = None, message: str | None = None) -> None:
|
||||
pytest.fail("Elicitation must not send unrelated progress")
|
||||
|
||||
|
||||
def prompt_session(channel: PromptChannel) -> ServerSession:
|
||||
connection: Final = Connection.for_loop(channel, protocol_version_hint="2025-11-25")
|
||||
connection.client_params = InitializeRequestParams(
|
||||
protocol_version="2025-11-25",
|
||||
capabilities=ClientCapabilities(elicitation=ElicitationCapability(form=FormElicitationCapability())),
|
||||
client_info=Implementation(name="component-client", version="1"),
|
||||
)
|
||||
return ServerSession(channel, connection)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_elicitation_is_blocked_without_affecting_ordinary_requests(
|
||||
monkeypatch: pytest.MonkeyPatch, registered_gateway: tuple[OperationContext, ObservedCalls, Sequence[str]]
|
||||
) -> None:
|
||||
context, _, _ = registered_gateway
|
||||
channel: Final = PromptChannel()
|
||||
session: Final = prompt_session(channel)
|
||||
session_token: Final = active_mcp_session_var.set(session)
|
||||
request_token: Final = active_mcp_request_ctx_var.set(
|
||||
ServerRequestContext(
|
||||
session=session,
|
||||
lifespan_context={},
|
||||
protocol_version="2025-11-25",
|
||||
method="tools/call",
|
||||
request_id="ordinary-request",
|
||||
)
|
||||
)
|
||||
params: Final = ElicitRequestFormParams(message="Confirm action", requested_schema={"type": "object"})
|
||||
ordinary: Final = create_elicitation_callback()
|
||||
|
||||
async def elicit() -> str:
|
||||
accepted: Final = await ordinary(None, params)
|
||||
assert isinstance(accepted, ElicitResult) and accepted.action == "accept"
|
||||
callback: Final = create_elicitation_callback()
|
||||
denied: Final = await asyncio.create_task(callback(None, params), context=Context())
|
||||
if isinstance(denied, ErrorData):
|
||||
raise MCPError(code=denied.code, message=denied.message)
|
||||
return denied.model_dump_json()
|
||||
|
||||
global_mcp_tool_registry.register_tool( # pyright: ignore[reportUnknownMemberType] # Legacy registry leaves handler Callable unparameterized.
|
||||
name="allowed-elicit",
|
||||
description="Confirm",
|
||||
input_schema={"type": "object"},
|
||||
handler=elicit,
|
||||
server_id="allowed",
|
||||
)
|
||||
global_mcp_server_manager.tool_name_to_mcp_server_name_mapping["allowed-elicit"] = "allowed"
|
||||
process: Final = install_transport(monkeypatch, (Call(id=0, tool_id="allowed-elicit", arguments={}),))
|
||||
try:
|
||||
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
|
||||
assert not result.is_error
|
||||
assert "unsupported in code mode. Call this tool directly" in error_text(process.replies[0])
|
||||
assert len(channel.prompts) == 1
|
||||
after: Final = await create_elicitation_callback()(None, params)
|
||||
assert isinstance(after, ElicitResult) and after.content == {"value": "accepted"}
|
||||
assert channel.prompts == (params.model_dump(mode="json", by_alias=True, exclude_none=True),) * 2
|
||||
finally:
|
||||
active_mcp_request_ctx_var.reset(request_token)
|
||||
active_mcp_session_var.reset(session_token)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError])
|
||||
async def test_code_context_resets_after_worker_failure_or_cancellation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
registered_gateway: tuple[OperationContext, ObservedCalls, Sequence[str]],
|
||||
failure: type[Exception] | type[asyncio.CancelledError],
|
||||
) -> None:
|
||||
context, _, _ = registered_gateway
|
||||
channel: Final = PromptChannel()
|
||||
token: Final = active_mcp_session_var.set(prompt_session(channel))
|
||||
|
||||
async def fail_start(command: Sequence[str], *, cwd: str, env: Mapping[str, str], stderr: int) -> Process:
|
||||
raise failure()
|
||||
|
||||
monkeypatch.setattr(anyio, "open_process", fail_start)
|
||||
try:
|
||||
if failure is asyncio.CancelledError:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await run_code_tool({"code": "/* protocol fixture */"}, context)
|
||||
else:
|
||||
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
|
||||
assert result.is_error
|
||||
response: Final = await create_elicitation_callback()(
|
||||
None, ElicitRequestFormParams(message="Still ordinary", requested_schema={"type": "object"})
|
||||
)
|
||||
assert isinstance(response, ElicitResult) and response.content == {"value": "accepted"}
|
||||
assert len(channel.prompts) == 1
|
||||
finally:
|
||||
active_mcp_session_var.reset(token)
|
||||
|
||||
|
||||
class CodeTextPolicy(CustomGuardrail):
|
||||
scanned: tuple[str, ...] = ()
|
||||
directions: tuple[str, ...] = ()
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: Mapping[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: LiteLLMLoggingObj | None = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.directions = (*self.directions, input_type)
|
||||
texts: Final = inputs.get("texts", [])
|
||||
self.scanned = (*self.scanned, *texts)
|
||||
if any("blocked-output" in text for text in texts):
|
||||
raise HTTPException(status_code=403, detail="Code output rejected")
|
||||
return {**inputs, "texts": [text.replace("mask-output", "filtered-output") for text in texts]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("nested", [False, True])
|
||||
@pytest.mark.parametrize("output", ["allowed-output", "mask-output", "blocked-output"])
|
||||
async def test_final_code_output_policy_enforces_and_preserves_logging(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
registered_gateway: tuple[OperationContext, ObservedCalls, Sequence[str]],
|
||||
nested: bool,
|
||||
output: str,
|
||||
) -> None:
|
||||
context, observed, received = registered_gateway
|
||||
policy: Final = CodeTextPolicy(
|
||||
guardrail_name="code-output", event_hook=GuardrailEventHooks.post_mcp_call, default_on=True
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, policy])
|
||||
calls: Final = (Call(id=0, tool_id="allowed-echo", arguments={"message": "safe-upstream"}),) if nested else ()
|
||||
process: Final = install_transport(monkeypatch, calls, (output,))
|
||||
if output == "blocked-output":
|
||||
with pytest.raises(HTTPException, match="Code output rejected"):
|
||||
await run_code_tool({"code": "/* protocol fixture */"}, context)
|
||||
assert len(observed.responses) == int(nested)
|
||||
install_transport(monkeypatch, (), ("allowed-output",))
|
||||
recovered: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
|
||||
assert not recovered.is_error
|
||||
else:
|
||||
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
|
||||
assert not result.is_error and isinstance(result.content[0], TextContent)
|
||||
assert load_json(result.content[0].text) == {
|
||||
"output": [output.replace("mask-output", "filtered-output")],
|
||||
"error": None,
|
||||
"calls_started": int(nested),
|
||||
}
|
||||
assert received == (["safe-upstream"] if nested else [])
|
||||
assert process.stopped
|
||||
assert len(observed.responses) == int(nested) + 1
|
||||
assert len(policy.scanned) == int(nested) + 1 + int(output == "blocked-output")
|
||||
assert policy.directions == ("response",) * len(policy.scanned)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_input_policy_blocks_before_worker_and_preserves_rewrites(
|
||||
monkeypatch: pytest.MonkeyPatch, registered_gateway: tuple[OperationContext, ObservedCalls, Sequence[str]]
|
||||
) -> None:
|
||||
context, observed, received = registered_gateway
|
||||
policy: Final = CodeTextPolicy(
|
||||
guardrail_name="code-input", event_hook=GuardrailEventHooks.pre_mcp_call, default_on=True
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, policy])
|
||||
process: Final = install_transport(monkeypatch, (), ("safe-output",))
|
||||
with pytest.raises(HTTPException, match="Code output rejected"):
|
||||
await run_code_tool({"code": "blocked-output"}, context)
|
||||
assert process.started is None and received == [] and observed.responses == ()
|
||||
result: Final = await run_code_tool({"code": "mask-output"}, context)
|
||||
assert not result.is_error
|
||||
assert process.started is not None and process.started.code == "filtered-output"
|
||||
assert policy.scanned == ("blocked-output", "mask-output")
|
||||
assert policy.directions == ("request", "request")
|
||||
assert len(observed.responses) == 1
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -44,9 +45,13 @@ def test_large_ids_as_strings_and_exact_safe_integers_survive() -> None:
|
|||
assert result.id == 7
|
||||
|
||||
|
||||
def test_pipe_limit_counts_unicode_bytes() -> None:
|
||||
@pytest.mark.parametrize(("value", "budget"), [("猫", 400), ("\ud800", 700)])
|
||||
def test_pipe_limit_counts_unicode_bytes(value: str, budget: int) -> None:
|
||||
reply: Final = Reply(id=0, value={"value": value * 100})
|
||||
encoded: Final = encode_message(reply, budget)
|
||||
assert Reply.model_validate(load_json(encoded.decode())) == reply
|
||||
with pytest.raises(ValueError, match="byte limit"):
|
||||
encode_message(Reply(id=0, value={"value": "猫" * 100}), 200)
|
||||
encode_message(reply, len(encoded) - 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.proxy.mcp_server.code_worker import (
|
|||
Reply,
|
||||
Start,
|
||||
encode_message,
|
||||
load_json,
|
||||
main,
|
||||
)
|
||||
|
||||
|
|
@ -26,7 +27,32 @@ def run_worker(monkeypatch: pytest.MonkeyPatch, request: bytes) -> list[Call | D
|
|||
monkeypatch.setattr(sys, "stdin", source)
|
||||
monkeypatch.setattr(sys, "stdout", destination)
|
||||
main()
|
||||
return [WORKER_MESSAGE.validate_json(line) for line in destination.buffer.getvalue().splitlines()]
|
||||
return [
|
||||
WORKER_MESSAGE.validate_python(load_json(line.decode())) for line in destination.buffer.getvalue().splitlines()
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("codepoint", "expected"), [(0xD800, "\ud800"), (0xDC00, "\udc00"), (0x1F680, "🚀"), (0x732B, "猫")]
|
||||
)
|
||||
def test_native_worker_preserves_unicode_output(monkeypatch: pytest.MonkeyPatch, codepoint: int, expected: str) -> None:
|
||||
limits: Final = CodeExecutionLimits()
|
||||
request: Final = Start(code=f"text('kept'); text(String.fromCodePoint({codepoint}));", limits=limits)
|
||||
assert run_worker(monkeypatch, encode_message(request, limits.message_bytes)) == [Done(output=("kept", expected))]
|
||||
|
||||
|
||||
def test_native_worker_roundtrips_surrogates_in_tool_arguments_and_results(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
limits: Final = CodeExecutionLimits()
|
||||
request: Final = Start(
|
||||
code="text('kept'); text(await tools.call_tool({tool_id:'echo',arguments:{value:String.fromCharCode(0xd800)}}));",
|
||||
limits=limits,
|
||||
)
|
||||
value: Final[dict[str, JsonValue]] = {"value": "\ud800"}
|
||||
messages: Final = run_worker(
|
||||
monkeypatch,
|
||||
encode_message(request, limits.message_bytes) + encode_message(Reply(id=0, value=value), limits.message_bytes),
|
||||
)
|
||||
assert messages == [Call(id=0, tool_id="echo", arguments=value), Done(output=("kept", value))]
|
||||
|
||||
|
||||
def test_native_worker_chains_full_tool_results(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue