69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
"""Redact credentials in MCP tool output before it reaches the model.
|
|
|
|
Covers every MCP route; in-process tools like ``Bash``/``Read`` are
|
|
guarded separately by ``agents.policy``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import mcp.types as mt
|
|
from fastmcp.server.middleware import Middleware
|
|
from fastmcp.tools.base import ToolResult
|
|
|
|
from beaver_gateway.security.redact import redact, redact_data
|
|
|
|
if TYPE_CHECKING:
|
|
from fastmcp.server.middleware import CallNext, MiddlewareContext
|
|
|
|
__all__ = ["RedactingMiddleware"]
|
|
|
|
|
|
def _redact_block(block: mt.ContentBlock) -> mt.ContentBlock:
|
|
"""Mask a content block's text; binary blocks pass through."""
|
|
if isinstance(block, mt.TextContent):
|
|
masked = redact(block.text)
|
|
return (
|
|
block if masked == block.text else block.model_copy(update={"text": masked})
|
|
)
|
|
if isinstance(block, mt.EmbeddedResource) and isinstance(
|
|
block.resource, mt.TextResourceContents
|
|
):
|
|
masked = redact(block.resource.text)
|
|
if masked == block.resource.text:
|
|
return block
|
|
resource = block.resource.model_copy(update={"text": masked})
|
|
return block.model_copy(update={"resource": resource})
|
|
return block
|
|
|
|
|
|
class RedactingMiddleware(Middleware):
|
|
"""Mask credentials in whatever a tool call returns.
|
|
|
|
Idempotent, so mounting a server both on its own namespace and in
|
|
the ``/mcp/all`` bundle costs a second pass, not a double mask.
|
|
"""
|
|
|
|
async def on_call_tool(
|
|
self,
|
|
context: MiddlewareContext[mt.CallToolRequestParams],
|
|
call_next: CallNext[mt.CallToolRequestParams, ToolResult],
|
|
) -> ToolResult:
|
|
"""Mask the result.
|
|
|
|
A non-``None`` ``meta`` on the masked copy takes the
|
|
``CallToolResult`` path that skips output-schema validation,
|
|
since masked structured content may no longer match it.
|
|
"""
|
|
result = await call_next(context)
|
|
content = [_redact_block(block) for block in result.content]
|
|
structured = redact_data(result.structured_content)
|
|
if content == result.content and structured == result.structured_content:
|
|
return result
|
|
return ToolResult(
|
|
content=content,
|
|
structured_content=structured,
|
|
meta=result.meta if result.meta is not None else {},
|
|
)
|