rsp2k
146bbbff25
soft_error_classifier gains error_keys, default ("error",). A dict result
whose value there is a non-empty string is a soft error, after the
status, flag and isError checks. None, blank strings, numbers, booleans,
dicts and lists do not count. error_keys=() turns it off.
401 lines
14 KiB
Python
401 lines
14 KiB
Python
"""Tests for per-tool exclusion, sampling and retention pruning.
|
|
|
|
Real FastMCP servers through fastmcp.Client, as in test_instrumentation.py.
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from fastmcp import Client, FastMCP
|
|
from fastmcp.exceptions import ToolError
|
|
from sqlalchemy import event, func, select
|
|
|
|
from fastmcp_feedback.instrumentation import (
|
|
DatabaseSink,
|
|
MemorySink,
|
|
ToolCallRecord,
|
|
instrument,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
async def db(db_url):
|
|
"""DatabaseSinks whose engines are always disposed, or the process hangs."""
|
|
made = []
|
|
|
|
def factory(name="calls", **kwargs):
|
|
sink = DatabaseSink(db_url(name), create_tables=True, **kwargs)
|
|
made.append(sink)
|
|
return sink
|
|
|
|
yield factory
|
|
for sink in made:
|
|
await sink.engine.dispose()
|
|
|
|
|
|
class Rolls:
|
|
"""Stand-in for random.random that returns the given values in turn."""
|
|
|
|
def __init__(self, *values):
|
|
self.values = list(values)
|
|
self.calls = 0
|
|
|
|
def __call__(self):
|
|
self.calls += 1
|
|
return self.values.pop(0)
|
|
|
|
|
|
def make_app(sinks=None, **kwargs):
|
|
sink = MemorySink()
|
|
app = FastMCP("sampling")
|
|
|
|
@app.tool
|
|
def poll() -> dict:
|
|
return {"status": "idle"}
|
|
|
|
@app.tool
|
|
def render(status: str) -> dict:
|
|
return {"status": status, "detail": f"render said {status}"}
|
|
|
|
@app.tool
|
|
def crash() -> dict:
|
|
raise RuntimeError("exploded")
|
|
|
|
@app.tool
|
|
def other() -> str:
|
|
return "fine"
|
|
|
|
mw = instrument(app, sinks if sinks is not None else [sink], **kwargs)
|
|
return app, sink, mw
|
|
|
|
|
|
async def call(app, mw, *calls):
|
|
"""Run (tool, args) pairs; return the results, None where the tool raised."""
|
|
results = []
|
|
async with Client(app) as client:
|
|
for tool, args in calls:
|
|
try:
|
|
results.append(await client.call_tool(tool, args))
|
|
except ToolError:
|
|
results.append(None)
|
|
await mw.flush()
|
|
return results
|
|
|
|
|
|
def old_record(tool="old", days=40, **kwargs):
|
|
return ToolCallRecord(
|
|
tool=tool,
|
|
started_at=datetime.now(UTC) - timedelta(days=days),
|
|
duration_ms=1.0,
|
|
ok=True,
|
|
mode="meta",
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
async def rows(sink, *columns):
|
|
t = sink.table
|
|
cols = [t.c[c] for c in columns] if columns else [t.c.tool]
|
|
async with sink.engine.connect() as conn:
|
|
return [
|
|
tuple(r)
|
|
for r in (await conn.execute(select(*cols).order_by(t.c.started_at))).all()
|
|
]
|
|
|
|
|
|
async def count(sink):
|
|
async with sink.engine.connect() as conn:
|
|
return (
|
|
await conn.execute(select(func.count()).select_from(sink.table))
|
|
).scalar()
|
|
|
|
|
|
class TestExclusion:
|
|
async def test_excluded_tool_not_recorded_nor_kept(self):
|
|
app, sink, mw = make_app(exclude_tools={"poll"})
|
|
results = await call(app, mw, ("poll", {}), ("other", {}), ("poll", {}))
|
|
assert [r.tool for r in sink.records] == ["other"]
|
|
assert [r.tool for r in mw.recent] == ["other"]
|
|
# The call itself is untouched.
|
|
assert results[0].structured_content == {"status": "idle"}
|
|
|
|
async def test_excluded_tool_failures_not_recorded_either(self):
|
|
app, sink, mw = make_app(exclude_tools=["crash"])
|
|
[result] = await call(app, mw, ("crash", {}))
|
|
assert result is None # the error still reached the client
|
|
assert sink.records == []
|
|
assert list(mw.recent) == []
|
|
|
|
async def test_excluded_tool_skips_hooks(self):
|
|
seen = []
|
|
app, sink, mw = make_app(
|
|
exclude_tools={"poll"}, identity_resolver=lambda c: seen.append(1) or {}
|
|
)
|
|
await call(app, mw, ("poll", {}))
|
|
assert seen == []
|
|
|
|
|
|
class TestSampling:
|
|
async def test_recorded_when_random_below_rate(self):
|
|
rolls = Rolls(0.29, 0.3, 0.31, 0.0)
|
|
app, sink, mw = make_app(sample_rates={"poll": 0.3}, random=rolls)
|
|
await call(app, mw, *[("poll", {})] * 4)
|
|
assert rolls.calls == 4
|
|
assert len(sink.records) == 2
|
|
assert [r.sample_rate for r in sink.records] == [0.3, 0.3]
|
|
assert all(r.outcome == "ok" for r in sink.records)
|
|
|
|
async def test_sampled_out_calls_not_in_recent(self):
|
|
app, sink, mw = make_app(sample_rates={"poll": 0.5}, random=Rolls(0.9, 0.1))
|
|
results = await call(app, mw, ("poll", {}), ("poll", {}))
|
|
assert len(sink.records) == 1
|
|
assert [r.id for r in mw.recent] == [sink.records[0].id]
|
|
# A sampled-out call still returns its result.
|
|
assert results[0].structured_content == {"status": "idle"}
|
|
|
|
async def test_sampled_out_ok_calls_skip_hooks(self):
|
|
seen = []
|
|
|
|
def who(context):
|
|
seen.append("identity")
|
|
return {}
|
|
|
|
def enrich(tool, args, result, context):
|
|
seen.append("enrich")
|
|
return {}
|
|
|
|
app, sink, mw = make_app(
|
|
sample_rates={"poll": 0.1},
|
|
random=Rolls(0.99),
|
|
identity_resolver=who,
|
|
enricher=enrich,
|
|
)
|
|
await call(app, mw, ("poll", {}))
|
|
assert sink.records == []
|
|
assert seen == []
|
|
|
|
async def test_unsampled_tools_recorded_without_rate_or_roll(self):
|
|
rolls = Rolls()
|
|
app, sink, mw = make_app(sample_rates={"poll": 0.1}, random=rolls)
|
|
await call(app, mw, ("other", {}), ("other", {}))
|
|
assert rolls.calls == 0
|
|
assert [(r.tool, r.sample_rate) for r in sink.records] == [
|
|
("other", None),
|
|
("other", None),
|
|
]
|
|
|
|
async def test_rate_of_one_records_everything_with_rate(self):
|
|
app, sink, mw = make_app(sample_rates={"poll": 1}, random=Rolls(0.999))
|
|
await call(app, mw, ("poll", {}))
|
|
[record] = sink.records
|
|
assert record.sample_rate == 1.0
|
|
|
|
async def test_exception_of_sampled_tool_always_recorded(self):
|
|
# Rolls would sample every call out; failures must not consult them.
|
|
rolls = Rolls(0.99, 0.99)
|
|
app, sink, mw = make_app(sample_rates={"crash": 0.01}, random=rolls)
|
|
await call(app, mw, ("crash", {}), ("crash", {}))
|
|
assert [(r.outcome, r.error_type, r.sample_rate) for r in sink.records] == [
|
|
("error", "RuntimeError", None),
|
|
("error", "RuntimeError", None),
|
|
]
|
|
assert len(mw.recent) == 2
|
|
|
|
async def test_soft_error_of_sampled_tool_always_recorded(self):
|
|
rolls = Rolls(0.99, 0.99)
|
|
app, sink, mw = make_app(sample_rates={"render": 0.01}, random=rolls)
|
|
await call(
|
|
app, mw, ("render", {"status": "failed"}), ("render", {"status": "done"})
|
|
)
|
|
assert [(r.outcome, r.sample_rate) for r in sink.records] == [
|
|
("soft_error", None)
|
|
]
|
|
assert sink.records[0].error_message == "status=failed: render said failed"
|
|
assert [r.id for r in mw.recent] == [sink.records[0].id]
|
|
|
|
async def test_sampled_in_failure_has_no_rate(self):
|
|
app, sink, mw = make_app(sample_rates={"crash": 0.5}, random=Rolls(0.0))
|
|
await call(app, mw, ("crash", {}))
|
|
[record] = sink.records
|
|
assert record.sample_rate is None
|
|
|
|
@pytest.mark.parametrize("rate", [0, -0.1, 1.01, 2, float("nan")])
|
|
def test_invalid_rates_rejected(self, rate):
|
|
with pytest.raises(ValueError, match="poll"):
|
|
make_app(sample_rates={"poll": rate})
|
|
|
|
def test_non_numeric_rate_rejected(self):
|
|
with pytest.raises((TypeError, ValueError)):
|
|
make_app(sample_rates={"poll": "often"})
|
|
|
|
async def test_estimated_count(self):
|
|
"""sum(1 / coalesce(sample_rate, 1)) estimates the true call count."""
|
|
rolls = Rolls(*[i / 10 for i in range(10)])
|
|
app, sink, mw = make_app(sample_rates={"poll": 0.2}, random=rolls)
|
|
await call(app, mw, *[("poll", {})] * 10, ("crash", {}))
|
|
estimate = sum(1 / (r.sample_rate or 1) for r in sink.records)
|
|
assert estimate == pytest.approx(11)
|
|
|
|
|
|
class TestDatabaseSinkColumn:
|
|
async def test_writes_sample_rate(self, db):
|
|
sink = db()
|
|
app, _, mw = make_app(
|
|
sinks=[sink],
|
|
sample_rates={"poll": 0.25, "crash": 0.25},
|
|
random=Rolls(0.1, 0.9),
|
|
)
|
|
await call(app, mw, ("poll", {}), ("poll", {}), ("crash", {}), ("other", {}))
|
|
assert await rows(sink, "tool", "outcome", "sample_rate") == [
|
|
("poll", "ok", 0.25),
|
|
("crash", "error", None),
|
|
("other", "ok", None),
|
|
]
|
|
|
|
|
|
class TestPrune:
|
|
async def test_deletes_old_rows_keeps_new(self, db):
|
|
sink = db()
|
|
await sink.write([old_record("a", days=40), old_record("b", days=31)])
|
|
await sink.write([old_record("c", days=29), old_record("d", days=0)])
|
|
deleted = await sink.prune(datetime.now(UTC) - timedelta(days=30))
|
|
assert deleted == 2
|
|
assert await rows(sink) == [("c",), ("d",)]
|
|
|
|
async def test_default_cutoff_uses_retention(self, db):
|
|
sink = db(retention=timedelta(days=30))
|
|
sink._last_prune = time.monotonic() # keep write() from pruning
|
|
await sink.write([old_record("a", days=40), old_record("b", days=1)])
|
|
assert await count(sink) == 2
|
|
assert await sink.prune() == 1
|
|
assert await rows(sink) == [("b",)]
|
|
|
|
async def test_no_retention_no_cutoff_is_a_no_op(self, db):
|
|
sink = db()
|
|
await sink.write([old_record(days=400)])
|
|
assert await sink.prune() == 0
|
|
assert await count(sink) == 1
|
|
|
|
async def test_keeps_linked_calls_and_links(self, db):
|
|
sink = db()
|
|
linked, plain = old_record("linked"), old_record("plain")
|
|
await sink.write([linked, plain])
|
|
await sink.replace_links(
|
|
"bug-1",
|
|
[
|
|
{
|
|
"feedback_ref": "bug-1",
|
|
"call_id": linked.id,
|
|
"position": 0,
|
|
"rule": "session",
|
|
"linked_at": datetime.now(UTC),
|
|
},
|
|
{
|
|
# A link to a call that is not in this database.
|
|
"feedback_ref": "bug-1",
|
|
"call_id": "missing",
|
|
"position": 1,
|
|
"rule": "session",
|
|
"linked_at": datetime.now(UTC),
|
|
},
|
|
],
|
|
)
|
|
assert await sink.prune(datetime.now(UTC)) == 1
|
|
assert await rows(sink) == [("linked",)]
|
|
context = await sink.linked_calls("bug-1")
|
|
assert [(r["call_id"], r["tool"]) for r in context] == [
|
|
(linked.id, "linked"),
|
|
("missing", None),
|
|
]
|
|
|
|
async def test_deletes_in_batches(self, db):
|
|
sink = db(prune_batch=3)
|
|
await sink.write(
|
|
[old_record(f"t{i}") for i in range(10)] + [old_record("new", 0)]
|
|
)
|
|
deletes = []
|
|
|
|
def spy(conn, cursor, statement, params, context, executemany):
|
|
if statement.lstrip().upper().startswith("DELETE"):
|
|
deletes.append(statement)
|
|
|
|
event.listen(sink.engine.sync_engine, "before_cursor_execute", spy)
|
|
try:
|
|
assert await sink.prune(datetime.now(UTC) - timedelta(days=1)) == 10
|
|
finally:
|
|
event.remove(sink.engine.sync_engine, "before_cursor_execute", spy)
|
|
assert len(deletes) == 4 # 3 + 3 + 3 + 1
|
|
assert await rows(sink) == [("new",)]
|
|
|
|
async def test_exact_multiple_of_batch(self, db):
|
|
sink = db(prune_batch=5)
|
|
await sink.write([old_record(f"t{i}") for i in range(10)])
|
|
assert await sink.prune(datetime.now(UTC)) == 10
|
|
assert await count(sink) == 0
|
|
|
|
def test_invalid_settings_rejected(self):
|
|
with pytest.raises(ValueError):
|
|
DatabaseSink("sqlite+aiosqlite://", retention=timedelta(0))
|
|
with pytest.raises(ValueError):
|
|
DatabaseSink("sqlite+aiosqlite://", prune_batch=0)
|
|
|
|
|
|
class TestAutomaticPrune:
|
|
async def test_first_write_prunes_then_waits_for_interval(self, db):
|
|
sink = db(retention=timedelta(days=30), prune_interval=timedelta(hours=1))
|
|
await sink.write([old_record("a"), old_record("new", 0)])
|
|
assert await rows(sink) == [("new",)]
|
|
|
|
await sink.write([old_record("b")])
|
|
assert await count(sink) == 2 # inside the interval: not pruned
|
|
|
|
sink._last_prune = time.monotonic() - 3601
|
|
await sink.write([old_record("c", 0)])
|
|
assert sorted(await rows(sink)) == [("c",), ("new",)]
|
|
|
|
async def test_no_retention_never_prunes(self, db):
|
|
sink = db()
|
|
await sink.write([old_record(days=4000)])
|
|
assert await count(sink) == 1
|
|
assert sink._last_prune is None
|
|
|
|
async def test_failing_prune_does_not_stop_inserts(self, db, caplog):
|
|
sink = db(retention=timedelta(days=30))
|
|
attempts = []
|
|
|
|
async def broken(older_than=None):
|
|
attempts.append(older_than)
|
|
raise RuntimeError("prune blew up")
|
|
|
|
sink.prune = broken
|
|
with caplog.at_level(logging.WARNING):
|
|
await sink.write([old_record("a", 0)])
|
|
await sink.write([old_record("b", 0)])
|
|
assert await rows(sink) == [("a",), ("b",)]
|
|
assert len(attempts) == 1 # not retried on every write
|
|
assert "prune blew up" in caplog.text
|
|
|
|
async def test_through_the_middleware(self, db):
|
|
sink = db(retention=timedelta(days=30))
|
|
sink._last_prune = time.monotonic()
|
|
await sink.write([old_record("stale")])
|
|
sink._last_prune = None
|
|
app, _, mw = make_app(sinks=[sink])
|
|
await call(app, mw, ("other", {}))
|
|
assert mw.dispatcher.sink_errors == 0
|
|
assert await rows(sink) == [("other",)]
|
|
|
|
async def test_prune_failure_through_middleware_keeps_records(self, db):
|
|
sink = db(retention=timedelta(days=30))
|
|
|
|
async def broken(older_than=None):
|
|
raise RuntimeError("no")
|
|
|
|
sink.prune = broken
|
|
app, _, mw = make_app(sinks=[sink])
|
|
await call(app, mw, ("other", {}), ("poll", {}))
|
|
assert mw.dispatcher.sink_errors == 0
|
|
assert sorted(await rows(sink)) == [("other",), ("poll",)]
|