86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
"""Redact credentials on the way out of an MCP tool.
|
|
|
|
A tool result is a wider channel than the log. It goes straight into the
|
|
model's context, from there into the turn record in Postgres, and from
|
|
there into the markdown transcript in ``💬 чаты`` — which Obsidian Sync
|
|
carries off the machine. One ``komodo`` deploy returns the resolved
|
|
compose file, ``environment:`` block and all, so a single call can put
|
|
every secret of a stack into all four places at once.
|
|
|
|
One filter, not a list of exceptions: every MCP the model can reach is
|
|
built into a ``FastMCP`` by :mod:`beaver_gateway.mcp.internal_app` —
|
|
``python_tool`` bundles like komodo, stdio subprocesses, remote HTTP
|
|
servers — and every route into one of them runs its middleware chain.
|
|
That covers the per-namespace ``/mcp/<name>/`` mounts claude-code talks
|
|
to, the ``/mcp/all`` bundle, the external MCP frontend reverse-proxying
|
|
into both, and the Raycast backend's direct ``call_tool``. A new MCP in
|
|
``config.py`` is covered the day it is added, without anyone
|
|
remembering to list it here.
|
|
|
|
What this does not reach: tools that never touch a FastMCP server — the
|
|
gateway's own ``gateway`` tools, and everything claude-code runs inside
|
|
its own process (``Bash``, ``Read``). Those are guarded by
|
|
:mod:`beaver_gateway.agents.policy` and the vault mounts instead.
|
|
"""
|
|
|
|
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:
|
|
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,
|
|
# Masked structured content no longer has to satisfy the
|
|
# tool's output schema; a non-None meta takes the
|
|
# ``CallToolResult`` path that skips that validation, the
|
|
# same trick ``ResponseLimitingMiddleware`` uses.
|
|
meta=result.meta if result.meta is not None else {},
|
|
)
|