fastmcp-feedback/tests/unit/test_instrumentation.py

"""Tests for per-call instrumentation: redaction, middleware, dispatch, sinks.

These run against real FastMCP servers through fastmcp.Client, so they
exercise the same middleware path on every supported FastMCP version.
"""

import asyncio
import io
import json
import time

import pytest
from fastmcp import Client, FastMCP
from fastmcp.exceptions import ToolError
from sqlalchemy import MetaData, inspect, select
from sqlalchemy.ext.asyncio import create_async_engine

from fastmcp_feedback import add_feedback_tools
from fastmcp_feedback.instrumentation import (
    REDACTED,
    CallbackSink,
    DatabaseSink,
    InstrumentationMiddleware,
    JsonLinesSink,
    MemorySink,
    RecordDispatcher,
    Redactor,
    ToolCallRecord,
    build_metadata,
    instrument,
)

BMCP_TOKEN = "bmcp_" + "A1b2C3d4E5f6G7h8I9j0"
JWT = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTYifQ.c2lnbmF0dXJlLXZhbHVl"


def make_app(sinks=None, **kwargs):
    sink = MemorySink()
    app = FastMCP("instrumented")

    @app.tool
    def echo(text: str) -> dict:
        return {"echo": text}

    @app.tool
    def create_token(label: str) -> dict:
        return {"label": label, "access_token": BMCP_TOKEN, "note": f"use {BMCP_TOKEN}"}

    @app.tool
    def login(user: str, password: str) -> str:
        return f"welcome {user}"

    @app.tool
    def fail(secret: str) -> str:
        raise ValueError(f"bad input secret={secret}")

    mw = instrument(app, sinks if sinks is not None else [sink], **kwargs)
    return app, sink, mw


async def call(app, tool, args=None):
    async with Client(app) as client:
        return await client.call_tool(tool, args or {})


class TestRedactor:
    def test_secret_keys_at_any_depth(self):
        data = {"a": {"b": [{"password": "hunter2", "keep": 1}]}, "Authorization": "x"}
        out = Redactor().redact(data)
        assert out["a"]["b"][0] == {"password": REDACTED, "keep": 1}
        assert out["Authorization"] == REDACTED

    def test_suffix_and_dash_variants(self):
        out = Redactor().redact(
            {
                "github_token": "t",
                "DB-Password": "p",
                "X-Api-Key": "k",
                "stripe_api_key": "s",
            }
        )
        assert set(out.values()) == {REDACTED}

    def test_non_secret_keys_untouched(self):
        data = {"max_tokens": 100, "token_count": 3, "tool": "echo", "tokens": [1, 2]}
        assert Redactor().redact(data) == data

    def test_secret_values_inside_text(self):
        text = f"pat {BMCP_TOKEN}, jwt {JWT}, header Bearer abcdefghijklmnop, key sk-{'x' * 20}"
        out = Redactor().redact_text(text)
        for secret in (BMCP_TOKEN, JWT, "abcdefghijklmnop", "sk-"):
            assert secret not in out
        assert out.count(REDACTED) == 4

    def test_key_value_text_keeps_the_key(self):
        out = Redactor().redact_text("failed: password=hunter2 api_key: abc123, ok")
        assert out == f"failed: password={REDACTED} api_key: {REDACTED}, ok"

    def test_input_not_mutated(self):
        data = {"password": "hunter2", "nested": {"token": "t"}}
        Redactor().redact(data)
        assert data == {"password": "hunter2", "nested": {"token": "t"}}

    def test_none_secret_value_left_as_none(self):
        assert Redactor().redact({"password": None}) == {"password": None}

    def test_extra_keys_and_patterns(self):
        r = Redactor(extra_keys=["license"], extra_patterns=[r"ACME-\d{6}"])
        assert r.redact({"license": "L", "msg": "code ACME-123456"}) == {
            "license": REDACTED,
            "msg": f"code {REDACTED}",
        }

    def test_depth_limit(self):
        deep = current = {}
        for _ in range(50):
            current["n"] = {}
            current = current["n"]
        assert "[TRUNCATED]" in json.dumps(Redactor(max_depth=5).redact(deep))


class TestMiddlewareRecording:
    async def test_success_in_meta_mode(self):
        app, sink, mw = make_app()
        await call(app, "echo", {"text": "hello"})
        await mw.flush()

        [rec] = sink.records
        assert rec.tool == "echo" and rec.ok and rec.mode == "meta"
        assert rec.duration_ms > 0
        assert rec.args_size == len(json.dumps({"text": "hello"}))
        assert rec.result_size and rec.result_size > 0
        assert rec.args is None and rec.result is None
        assert rec.session_id and rec.request_id

    async def test_full_mode_redacts_args_and_results(self):
        app, sink, mw = make_app(mode="full")
        await call(app, "create_token", {"label": "ci"})
        await call(app, "login", {"user": "ryan", "password": "hunter2"})
        await mw.flush()

        token_rec, login_rec = sink.records
        dumped = json.dumps([token_rec.to_dict(), login_rec.to_dict()])
        assert BMCP_TOKEN not in dumped and "hunter2" not in dumped
        assert token_rec.result["access_token"] == REDACTED
        assert token_rec.result["note"] == f"use {REDACTED}"
        assert token_rec.result["label"] == "ci"
        assert login_rec.args == {"user": "ryan", "password": REDACTED}

    async def test_failed_call_recorded_and_error_unchanged(self):
        app, sink, mw = make_app()
        with pytest.raises(ToolError, match="bad input"):
            await call(app, "fail", {"secret": "hunter2"})
        await mw.flush()

        [rec] = sink.records
        assert not rec.ok
        assert rec.error_type == "ValueError"
        assert "hunter2" not in rec.error_message
        assert f"secret={REDACTED}" in rec.error_message

    async def test_meta_only_tools_override_full_mode(self):
        app, sink, mw = make_app(mode="full", meta_only_tools={"create_token"})
        await call(app, "create_token", {"label": "ci"})
        await call(app, "echo", {"text": "hi"})
        await mw.flush()

        token_rec, echo_rec = sink.records
        assert token_rec.mode == "meta" and token_rec.result is None
        assert token_rec.result_size > 0
        assert echo_rec.mode == "full" and echo_rec.result == {"echo": "hi"}

    async def test_off_mode_records_nothing(self):
        app, sink, mw = make_app(mode="off")
        await call(app, "echo", {"text": "hi"})
        await mw.flush()
        assert sink.records == []

    async def test_mode_from_environment(self, monkeypatch):
        monkeypatch.setenv("FEEDBACK_INSTRUMENTATION_MODE", "full")
        assert InstrumentationMiddleware([]).mode == "full"

    def test_invalid_mode_rejected(self):
        with pytest.raises(ValueError, match="mode must be one of"):
            InstrumentationMiddleware([], mode="everything")


class TestHooks:
    async def test_identity_resolver_fills_columns(self):
        def who(context):
            assert context.fastmcp_context is not None
            return {
                "user_sub": "user-1",
                "caller_kind": "pat:7",
                "client_id": "blender-addon",
                "org": "acme",
                "token": BMCP_TOKEN,
            }

        app, sink, mw = make_app(identity_resolver=who)
        await call(app, "echo", {"text": "hi"})
        await mw.flush()

        [rec] = sink.records
        assert (rec.user_sub, rec.caller_kind, rec.client_id) == (
            "user-1",
            "pat:7",
            "blender-addon",
        )
        assert rec.identity == {"org": "acme", "token": REDACTED}

    async def test_async_enricher_gets_result_and_sets_extra(self):
        seen = {}

        async def enrich(tool, args, result, context):
            seen.update(tool=tool, args=args, result=result)
            return {"job_id": "job-42", "server_version": "2026.9.27", "api_key": "k"}

        app, sink, mw = make_app(enricher=enrich, server_version="default")
        await call(app, "echo", {"text": "hi"})
        await mw.flush()

        [rec] = sink.records
        assert seen == {
            "tool": "echo",
            "args": {"text": "hi"},
            "result": {"echo": "hi"},
        }
        assert rec.extra == {"job_id": "job-42", "api_key": REDACTED}
        assert rec.server_version == "2026.9.27"

    async def test_server_version_default(self):
        app, sink, mw = make_app(server_version="1.0")
        await call(app, "echo", {"text": "hi"})
        await mw.flush()
        assert sink.records[0].server_version == "1.0"

    async def test_hook_exceptions_do_not_break_calls(self):
        def boom(*args):
            raise RuntimeError("hook bug")

        app, sink, mw = make_app(identity_resolver=boom, enricher=boom)
        result = await call(app, "echo", {"text": "hi"})
        await mw.flush()

        assert result.data == {"echo": "hi"}
        [rec] = sink.records
        assert rec.ok and rec.identity == {} and rec.extra == {}

    async def test_slow_async_hook_is_cut_off(self):
        async def slow(*args):
            await asyncio.sleep(5)
            return {"never": True}

        app, sink, mw = make_app(enricher=slow, hook_timeout=0.05)
        start = time.perf_counter()
        await call(app, "echo", {"text": "hi"})
        assert time.perf_counter() - start < 2
        await mw.flush()
        assert sink.records[0].extra == {}


class TestNeverBlocksOrBreaks:
    async def test_failing_sink_does_not_break_calls_or_other_sinks(self):
        class Broken:
            async def write(self, records):
                raise ConnectionError("database down")

        good = MemorySink()
        app, _, mw = make_app(sinks=[Broken(), good])
        result = await call(app, "echo", {"text": "hi"})
        await mw.flush()

        assert result.data == {"echo": "hi"}
        assert mw.dispatcher.sink_errors == 1
        assert len(good.records) == 1

    async def test_stalled_sink_drops_records_instead_of_waiting(self):
        release = asyncio.Event()

        class Stalled:
            def __init__(self):
                self.records = []

            async def write(self, records):
                await release.wait()
                self.records.extend(records)

        stalled = Stalled()
        app, _, mw = make_app(sinks=[stalled], max_queue=2, batch_size=1)

        start = time.perf_counter()
        async with Client(app) as client:
            for i in range(6):
                await client.call_tool("echo", {"text": str(i)})
        assert time.perf_counter() - start < 3

        # One record is in the stalled write, two fit the queue, the rest drop.
        assert mw.dispatcher.dropped == 3
        release.set()
        await mw.flush()
        assert len(stalled.records) == 3

    async def test_record_building_failure_does_not_break_calls(self, monkeypatch):
        app, sink, mw = make_app()

        async def explode(*args, **kwargs):
            raise RuntimeError("bug in record building")

        monkeypatch.setattr(mw, "_build_record", explode)
        result = await call(app, "echo", {"text": "hi"})
        assert result.data == {"echo": "hi"}

    async def test_tool_error_passes_through_even_if_recording_fails(self, monkeypatch):
        app, sink, mw = make_app()

        async def explode(*args, **kwargs):
            raise RuntimeError("bug in record building")

        monkeypatch.setattr(mw, "_build_record", explode)
        with pytest.raises(ToolError, match="bad input"):
            await call(app, "fail", {"secret": "x"})

    def test_submit_without_event_loop_drops(self):
        dispatcher = RecordDispatcher([MemorySink()])
        record = ToolCallRecord(
            tool="t",
            started_at=__import__("datetime").datetime.now(),
            duration_ms=1,
            ok=True,
            mode="meta",
        )
        assert dispatcher.submit(record) is False
        assert dispatcher.dropped == 1


class TestSinks:
    async def test_json_lines_sink(self):
        stream = io.StringIO()
        app, _, mw = make_app(sinks=[JsonLinesSink(stream)])
        await call(app, "echo", {"text": "hi"})
        await call(app, "echo", {"text": "again"})
        await mw.flush()

        lines = [json.loads(line) for line in stream.getvalue().splitlines()]
        assert [line["tool"] for line in lines] == ["echo", "echo"]
        assert lines[0]["started_at"].endswith("+00:00")

    async def test_callback_sink_sync_and_async(self):
        got_sync, got_async = [], []

        async def async_cb(records):
            got_async.extend(records)

        app, _, mw = make_app(
            sinks=[CallbackSink(got_sync.extend), CallbackSink(async_cb)]
        )
        await call(app, "echo", {"text": "hi"})
        await mw.flush()
        assert len(got_sync) == len(got_async) == 1

    async def test_aclose_closes_sinks(self):
        closed = []

        class Closable(MemorySink):
            async def aclose(self):
                closed.append(True)

        app, _, mw = make_app(sinks=[Closable()])
        await call(app, "echo", {"text": "hi"})
        await mw.aclose()
        assert closed == [True]


class TestDatabaseSink:
    async def test_end_to_end_with_created_tables(self, tmp_path):
        url = f"sqlite+aiosqlite:///{tmp_path}/calls.db"
        db = DatabaseSink(url, create_tables=True)
        app, _, mw = make_app(sinks=[db], mode="full", server_version="9.9")
        await call(app, "create_token", {"label": "ci"})
        with pytest.raises(ToolError):
            await call(app, "fail", {"secret": "hunter2"})
        await mw.flush()

        async with db.engine.connect() as conn:
            rows = (
                (await conn.execute(select(db.table).order_by(db.table.c.started_at)))
                .mappings()
                .all()
            )
        await mw.aclose()

        assert [r["tool"] for r in rows] == ["create_token", "fail"]
        assert rows[0]["ok"] is True and rows[1]["ok"] is False
        assert rows[0]["result"]["access_token"] == REDACTED
        assert rows[1]["error_type"] == "ValueError"
        assert rows[0]["server_version"] == "9.9"
        assert "hunter2" not in json.dumps([dict(r) for r in rows], default=str)

    async def test_prefix_and_host_owned_engine(self, tmp_path):
        engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path}/host.db")
        db = DatabaseSink(engine, prefix="qa_", create_tables=True)
        app, _, mw = make_app(sinks=[db])
        await call(app, "echo", {"text": "hi"})
        await mw.aclose()

        # The host's engine is still usable: the sink did not dispose it.
        async with engine.connect() as conn:
            tables = await conn.run_sync(lambda c: inspect(c).get_table_names())
        await engine.dispose()
        assert sorted(tables) == [
            "qa_events",
            "qa_feedback_call_links",
            "qa_tool_calls",
        ]

    async def test_no_tables_created_by_default(self, tmp_path):
        engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path}/host.db")
        app, _, mw = make_app(sinks=[DatabaseSink(engine)])
        result = await call(app, "echo", {"text": "hi"})
        await mw.flush()

        async with engine.connect() as conn:
            tables = await conn.run_sync(lambda c: inspect(c).get_table_names())
        await engine.dispose()

        # The insert fails (no table), is counted, and the call still worked.
        assert result.data == {"echo": "hi"}
        assert tables == []
        assert mw.dispatcher.sink_errors == 1

    async def test_absent_payloads_are_sql_null(self, tmp_path):
        """Meta mode stores no payloads; they must be SQL NULL, not JSON null."""
        db = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/n.db", create_tables=True)
        app, _, mw = make_app(sinks=[db])
        await call(app, "echo", {"text": "hi"})
        await mw.flush()

        t = db.table
        async with db.engine.connect() as conn:
            nulls = (
                await conn.execute(select(t.c.args.is_(None), t.c.result.is_(None)))
            ).one()
        await mw.aclose()
        assert tuple(nulls) == (True, True)

    def test_build_metadata_joins_host_metadata(self):
        host = MetaData()
        assert build_metadata("ffb_", host) is host
        build_metadata("ffb_", host)
        assert sorted(host.tables) == [
            "ffb_events",
            "ffb_feedback_call_links",
            "ffb_tool_calls",
        ]


class TestIntegration:
    async def test_records_mounted_server_tools_by_prefixed_name(self):
        main = FastMCP("main")
        sub = FastMCP("sub")

        @sub.tool
        def ping() -> str:
            return "pong"

        main.mount(sub, "svc")
        sink = MemorySink()
        mw = instrument(main, [sink])
        await call(main, "svc_ping")
        await mw.flush()
        assert [r.tool for r in sink.records] == ["svc_ping"]

    async def test_records_feedback_tools(self):
        app = FastMCP("with-feedback")
        add_feedback_tools(app, database_url="sqlite:///:memory:")
        sink = MemorySink()
        mw = instrument(app, [sink])
        await call(
            app,
            "submit_feedback",
            {
                "request": {
                    "type": "bug",
                    "title": "t",
                    "description": "instrumented",
                    "submitter": "me",
                }
            },
        )
        await mw.flush()
        assert [r.tool for r in sink.records] == ["submit_feedback"]