beaver-gateway/tests/test_mcp_redacting.py

from __future__ import annotations

import mcp.types as mt
import pytest
from fastmcp import FastMCP
from fastmcp.tools.base import ToolResult

from beaver_gateway.core import redact as redact_mod
from beaver_gateway.core.redact import redact
from beaver_gateway.mcp.internal_app import build_internal_app
from beaver_gateway.mcp.redacting import RedactingMiddleware
from beaver_gateway.mcp.types import McpServer

# A trimmed but otherwise verbatim ``komodo deploy`` result: this is the
# shape Komodo returns for the ``Compose Config`` stage of a
# DeployStackService, which the komodo tool passes on to the model as
# ``-- <stage> ok`` plus the stage's stdout. Values are substituted.
KOMODO_DEPLOY = """DeployStackService Complete ok
-- Service/s ok
gateway
-- Compose Config ok
name: beaver-agent
services:
  gateway:
    container_name: beaver-gateway
    environment:
      ADMIN_PASS: 0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f50
      ADMIN_USER: admin
      BEAVERGRAM_MCP: https://tg.example.dev/mcp?token=bg-0123456789abcdef
      BOOTSTRAP_TOKENS: admin:bs-0123456789abcdef,komodo:hook-0123456789
      CALENDAR_MCPS: personal=https://cal.example.com/api/mcp?icsUrl=https%3A%2F%2Fcalendar.google.com%2Fcalendar%2Fical%2Fx%2Fprivate-0123456789abcdef%2Fbasic.ics
      CLAUDE_CODE_OAUTH_TOKEN: sk-ant-oat01-0123456789abcdef
      CLAUDE_HOME: /home/beaver-runner
      COMPOSE_PROFILES: db,gateway,obsidian,t3
      DATABASE_URL: postgresql+psycopg://beaver:pg-0123456789abcdef@postgres:5432/beaver
      KOMODO_URL: https://komod.example.dev
      POSTGRES_DB: beaver
      POSTGRES_PASSWORD: pg-0123456789abcdef
      POSTGRES_USER: beaver
    ports:
      - 62990:62990
-- Compose Up ok
 Container beaver-gateway  Started
"""

SECRETS = (
    "0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f50",
    "bg-0123456789abcdef",
    "bs-0123456789abcdef",
    "hook-0123456789",
    "private-0123456789abcdef",
    "sk-ant-oat01-0123456789abcdef",
    "pg-0123456789abcdef",
)


@pytest.fixture(autouse=True)
def _no_env_secrets() -> None:
    """Prove the patterns alone carry it — no literal env values loaded."""
    redact_mod._ENV_SECRETS[:] = []  # noqa: SLF001


def test_compose_config_dump_loses_every_value_but_keeps_every_name() -> None:
    out = redact(KOMODO_DEPLOY)
    for secret in SECRETS:
        assert secret not in out, secret
    # Names stay: the dump still says what is configured.
    for name in (
        "ADMIN_PASS",
        "BOOTSTRAP_TOKENS",
        "CALENDAR_MCPS",
        "POSTGRES_PASSWORD",
        # a username is half of a login pair — in prod ADMIN_USER is a
        # 64-hex string, not "admin"
        "ADMIN_USER",
        "POSTGRES_USER",
    ):
        assert f"{name}: <…>" in out
    # Non-secret values are untouched, so the result stays worth reading.
    assert "POSTGRES_DB: beaver" in out
    assert "CLAUDE_HOME: /home/beaver-runner" in out
    assert "COMPOSE_PROFILES: db,gateway,obsidian,t3" in out
    assert "KOMODO_URL: https://komod.example.dev" in out
    assert "container_name: beaver-gateway" in out
    assert "-- Compose Up ok" in out
    assert " Container beaver-gateway  Started" in out


def test_database_url_keeps_its_shape_and_loses_its_password() -> None:
    # DATABASE_URL is not a secret-looking name, so the whole line is not
    # masked; the password is caught as the userinfo of the connection
    # string, which is what keeps the host and database readable.
    out = redact(KOMODO_DEPLOY)
    line = next(ln for ln in out.splitlines() if "DATABASE_URL" in ln)
    assert "pg-0123456789abcdef" not in line
    assert "postgres:5432/beaver" in line


def test_redaction_is_idempotent() -> None:
    once = redact(KOMODO_DEPLOY)
    assert redact(once) == once


def test_a_log_line_is_not_swallowed_by_the_assignment_rule() -> None:
    # ``token:`` in mid-line must not eat the rest of an audit line.
    line = "mcp: actor=token:komodo namespace=calendar POST /mcp/ -> 200"
    assert redact(line) == line


@pytest.mark.anyio
async def test_middleware_masks_a_tool_result() -> None:
    def deploy(stack: str) -> str:
        """Deploy a stack."""
        return f"{stack}\n{KOMODO_DEPLOY}"

    server = FastMCP(name="komodo", tools=[deploy])
    server.add_middleware(RedactingMiddleware())
    result = await server.call_tool("deploy", {"stack": "beaver-agent"})
    text = "\n".join(b.text for b in result.content if isinstance(b, mt.TextContent))
    for secret in SECRETS:
        assert secret not in text, secret
    assert "ADMIN_PASS: <…>" in text


@pytest.mark.anyio
async def test_middleware_masks_structured_content() -> None:
    def config() -> dict[str, object]:
        """Return a config blob."""
        return {
            "env": {"ADMIN_PASS": "0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f50"},
            "port": 62990,
        }

    server = FastMCP(name="cfg", tools=[config])
    server.add_middleware(RedactingMiddleware())
    result = await server.call_tool("config", {})
    dumped = str(result.structured_content) + str(result.content)
    assert "0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f50" not in dumped
    assert "62990" in dumped


@pytest.mark.anyio
async def test_clean_result_passes_through_untouched() -> None:
    def ping() -> str:
        """Say hello."""
        return "pong"

    server = FastMCP(name="p", tools=[ping])
    server.add_middleware(RedactingMiddleware())
    result = await server.call_tool("ping", {})
    assert isinstance(result, ToolResult)
    assert result.content[0].text == "pong"  # ty: ignore[unresolved-attribute]


@pytest.mark.anyio
async def test_every_namespace_built_by_the_aggregator_is_filtered() -> None:
    def deploy() -> str:
        """Deploy."""
        return KOMODO_DEPLOY

    _app, _urls, servers = build_internal_app(
        [McpServer.python_tool(name="komodo", tools=[deploy])],
        host="127.0.0.1",
        port=8765,
    )
    for name, server in servers.items():
        result = await server.call_tool(
            "deploy" if name == "komodo" else "komodo_deploy", {}
        )
        text = "\n".join(
            b.text for b in result.content if isinstance(b, mt.TextContent)
        )
        assert "0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f50" not in text, name


@pytest.mark.anyio
async def test_gateway_own_tools_are_filtered_too() -> None:
    # These bypass the FastMCP middleware — the SDK mounts them
    # in-process — so they carry their own wrapper.
    from claude_agent_sdk import SdkMcpTool

    from beaver_gateway.core.gateway_tools import _redacting

    async def handler(_args: dict[str, object]) -> dict[str, object]:
        return {"content": [{"type": "text", "text": KOMODO_DEPLOY}]}

    spec = _redacting(
        SdkMcpTool(
            name="read_conversation", description="", input_schema={}, handler=handler
        )
    )
    out = str(await spec.handler({}))
    for secret in SECRETS:
        assert secret not in out, secret


def test_source_code_the_dispatcher_reads_is_not_mangled() -> None:
    # t3code hands back files and diffs. A SCREAMING_SNAKE constant is
    # an assignment too, but ``NAME = expr`` is code, not a config dump.
    source = (
        'SESSION_SECRET = os.environ["SESSION_SECRET"]\nAPI_KEY_HEADER = "X-Api-Key"'
    )
    assert redact(source) == source


def test_env_file_lines_are_still_masked() -> None:
    out = redact("ADMIN_PASS=0d1e2f3a4b5c6d7e8f9a0b1c2d3e4f50\nPOSTGRES_DB=beaver")
    assert out == "ADMIN_PASS=<…>\nPOSTGRES_DB=beaver"