From 55eeadafa7af58d28d2cb0bb86e1ba599524f5ac Mon Sep 17 00:00:00 2001 From: moe-berri <322602897+moe-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:37:25 -0700 Subject: [PATCH] fix(mcp): enforce code output policies and interaction boundaries --- litellm/proxy/mcp_server/code_execution.py | 14 +- litellm/proxy/mcp_server/code_runtime.py | 4 +- litellm/proxy/mcp_server/code_worker.py | 13 +- litellm/proxy/mcp_server/legacy_callbacks.py | 24 +- tests/integration/mcp/test_interactions.py | 162 ++++++++++++ .../integration/mcp/test_mcp_code_runtime.py | 16 ++ .../mcp_server/test_code_orchestration.py | 238 +++++++++++++++++- .../proxy/mcp_server/test_code_runtime.py | 9 +- .../unit/proxy/mcp_server/test_code_worker.py | 28 ++- 9 files changed, 488 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/mcp_server/code_execution.py b/litellm/proxy/mcp_server/code_execution.py index 0cc4b4f58cd..0e710bdd735 100644 --- a/litellm/proxy/mcp_server/code_execution.py +++ b/litellm/proxy/mcp_server/code_execution.py @@ -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( diff --git a/litellm/proxy/mcp_server/code_runtime.py b/litellm/proxy/mcp_server/code_runtime.py index 5e1d5353b4d..fe9359f85fc 100644 --- a/litellm/proxy/mcp_server/code_runtime.py +++ b/litellm/proxy/mcp_server/code_runtime.py @@ -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) ) diff --git a/litellm/proxy/mcp_server/code_worker.py b/litellm/proxy/mcp_server/code_worker.py index 78b3f3b9d4e..6785685889b 100644 --- a/litellm/proxy/mcp_server/code_worker.py +++ b/litellm/proxy/mcp_server/code_worker.py @@ -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 diff --git a/litellm/proxy/mcp_server/legacy_callbacks.py b/litellm/proxy/mcp_server/legacy_callbacks.py index 3d2f97492bb..0ce012a6071 100644 --- a/litellm/proxy/mcp_server/legacy_callbacks.py +++ b/litellm/proxy/mcp_server/legacy_callbacks.py @@ -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, diff --git a/tests/integration/mcp/test_interactions.py b/tests/integration/mcp/test_interactions.py index 5577cf16067..f24a0745e49 100644 --- a/tests/integration/mcp/test_interactions.py +++ b/tests/integration/mcp/test_interactions.py @@ -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" diff --git a/tests/integration/mcp/test_mcp_code_runtime.py b/tests/integration/mcp/test_mcp_code_runtime.py index 2f429efe383..d887b316ec9 100644 --- a/tests/integration/mcp/test_mcp_code_runtime.py +++ b/tests/integration/mcp/test_mcp_code_runtime.py @@ -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)" diff --git a/tests/unit/proxy/mcp_server/test_code_orchestration.py b/tests/unit/proxy/mcp_server/test_code_orchestration.py index 34ff34ef1b3..daaf1570a0b 100644 --- a/tests/unit/proxy/mcp_server/test_code_orchestration.py +++ b/tests/unit/proxy/mcp_server/test_code_orchestration.py @@ -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 diff --git a/tests/unit/proxy/mcp_server/test_code_runtime.py b/tests/unit/proxy/mcp_server/test_code_runtime.py index 1e6b05d6ce4..619212944c6 100644 --- a/tests/unit/proxy/mcp_server/test_code_runtime.py +++ b/tests/unit/proxy/mcp_server/test_code_runtime.py @@ -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 diff --git a/tests/unit/proxy/mcp_server/test_code_worker.py b/tests/unit/proxy/mcp_server/test_code_worker.py index ff7ba3285b5..063ce26722c 100644 --- a/tests/unit/proxy/mcp_server/test_code_worker.py +++ b/tests/unit/proxy/mcp_server/test_code_worker.py @@ -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: