diff --git a/litellm/proxy/mcp_server/code_execution.py b/litellm/proxy/mcp_server/code_execution.py index 0e710bdd735..cf8c3b187ff 100644 --- a/litellm/proxy/mcp_server/code_execution.py +++ b/litellm/proxy/mcp_server/code_execution.py @@ -281,6 +281,7 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon if caller is None: return error_result("MCP code execution requires an authenticated caller") from litellm.proxy.mcp_server.code_runtime import execute_code + from litellm.proxy.mcp_server.code_worker import json_bytes try: request: Final = CodeInput.model_validate(arguments or {}) @@ -375,9 +376,9 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon content=[ TextContent( type="text", - text=json.dumps( + text=json_bytes( {"output": result.output, "error": result.error, "calls_started": result.calls_started} - ), + ).decode("utf-8"), ) ], is_error=result.error is not None, diff --git a/tests/integration/mcp/test_interactions.py b/tests/integration/mcp/test_interactions.py index f24a0745e49..cff3e927d7d 100644 --- a/tests/integration/mcp/test_interactions.py +++ b/tests/integration/mcp/test_interactions.py @@ -1702,7 +1702,7 @@ 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", [])): + if any("qa-rejected-output" in text or "秘密" in text for text in inputs.get("texts", [])): raise HTTPException(status_code=400, detail="QA output rejected") return inputs """) @@ -1746,7 +1746,12 @@ class RejectOutput(CustomGuardrail): outcomes = [] for path in ("/mcp", "/mcp/proxy", "rest"): responses = [] - for code in ('text("qa-rejected-output")', 'text("qa-allowed-output")'): + for code in ( + 'text("qa-rejected-output")', + 'text(String.fromCodePoint(0x79d8, 0x5bc6))', + 'text("qa-allowed-output")', + 'text(String.fromCodePoint(0x4f60, 0x597d))', + ): if path == "rest": result = McpCaller(gateway, gateway.key, "rest").call("litellm_execute_code", {"code": code}) else: @@ -1755,11 +1760,18 @@ class RejectOutput(CustomGuardrail): ) 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) + unicode_outcomes: Final = tuple( + {"route": route, "status": result.status, "ok": result.ok, "output": result.text} + for route, _rejected, result, _allowed, _unicode_allowed in outcomes + ) + for path, rejected, unicode_rejected, allowed, unicode_allowed in outcomes: + for blocked in (rejected, unicode_rejected): + assert not blocked.ok, (path, blocked, unicode_outcomes) + assert "QA output rejected" in (blocked.error or ""), (path, blocked) assert allowed.ok and allowed.text is not None, (path, allowed) assert json.loads(allowed.text)["output"] == ["qa-allowed-output"] + assert unicode_allowed.ok and unicode_allowed.text is not None, (path, unicode_allowed) + assert json.loads(unicode_allowed.text)["output"] == ["你好"] surrogate = McpCaller(gateway, gateway.key, "rest").call( "litellm_execute_code", {"code": "text('kept'); text(String.fromCharCode(0xd800));"} ) diff --git a/tests/unit/proxy/mcp_server/test_code_orchestration.py b/tests/unit/proxy/mcp_server/test_code_orchestration.py index daaf1570a0b..011279a01ba 100644 --- a/tests/unit/proxy/mcp_server/test_code_orchestration.py +++ b/tests/unit/proxy/mcp_server/test_code_orchestration.py @@ -578,14 +578,29 @@ class CodeTextPolicy(CustomGuardrail): 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): + if any("blocked-output" in text or "秘密" 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]} + return { + **inputs, + "texts": [ + text.replace("mask-output", "filtered-output").replace("机密", "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"]) +@pytest.mark.parametrize( + "output", + [ + "allowed-output", + "mask-output", + "blocked-output", + pytest.param("秘密", id="unicode-block"), + pytest.param("机密", id="unicode-mask"), + pytest.param("你好", id="unicode-allow"), + ], +) async def test_final_code_output_policy_enforces_and_preserves_logging( monkeypatch: pytest.MonkeyPatch, registered_gateway: tuple[OperationContext, ObservedCalls, Sequence[str]], @@ -599,7 +614,7 @@ async def test_final_code_output_policy_enforces_and_preserves_logging( 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": + if output in ("blocked-output", "秘密"): with pytest.raises(HTTPException, match="Code output rejected"): await run_code_tool({"code": "/* protocol fixture */"}, context) assert len(observed.responses) == int(nested) @@ -610,15 +625,16 @@ async def test_final_code_output_policy_enforces_and_preserves_logging( 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")], + "output": [output.replace("mask-output", "filtered-output").replace("机密", "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 len(policy.scanned) == int(nested) + 1 + int(output in ("blocked-output", "秘密")) assert policy.directions == ("response",) * len(policy.scanned) + assert any(output in text for text in policy.scanned) @pytest.mark.asyncio