rsp2k
2e1278d81f
InstrumentationMiddleware (and instrument(app)) records every tool call on a FastMCP server: tool, timing, outcome and original error type, session and request ids, argument and result sizes, plus identity and enrichment hook output. Modes off/meta/full; full adds redacted args and results. Recording can't slow or break a call. Records reach sinks through a bounded queue drained by a background task and are dropped, counted, when it fills. Sink and hook failures are logged, never raised, and async hooks are cut off after hook_timeout. Redactor strips secrets by key at any depth and by value pattern from args, results, error text and hook output. DatabaseSink runs on async SQLAlchemy, takes a host AsyncEngine, writes to a prefixed table and never creates it unless asked; build_metadata() plugs into Alembic env.py. Verified on FastMCP 2.12.2, 2.14.3, 3.3.1, 3.4.7 and 4.0.10. CI matrix gains 3.x. Version 2026.09.27.1.
118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
"""Background delivery of records to sinks, off the tool-call path."""
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from collections.abc import Sequence
|
|
|
|
from .record import ToolCallRecord
|
|
from .sinks import Sink
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RecordDispatcher:
|
|
"""Queue records and hand them to sinks from a background task.
|
|
|
|
``submit`` never blocks and never raises: when the queue is full the record
|
|
is dropped and counted, so a slow or broken sink cannot slow down or fail
|
|
a tool call. Sink errors are logged and counted, never propagated.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sinks: Sequence[Sink],
|
|
*,
|
|
max_queue: int = 1000,
|
|
batch_size: int = 100,
|
|
sink_timeout: float = 10.0,
|
|
):
|
|
self.sinks = list(sinks)
|
|
self.max_queue = max_queue
|
|
self.batch_size = batch_size
|
|
self.sink_timeout = sink_timeout
|
|
self.dropped = 0
|
|
self.sink_errors = 0
|
|
self._queue: asyncio.Queue[ToolCallRecord] | None = None
|
|
self._task: asyncio.Task | None = None
|
|
self._loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
def submit(self, record: ToolCallRecord) -> bool:
|
|
"""Queue a record for delivery. Returns False if it was dropped."""
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
self._drop("no running event loop")
|
|
return False
|
|
|
|
# A queue belongs to one loop; tests and some hosts run several.
|
|
if self._queue is None or self._loop is not loop:
|
|
self._loop = loop
|
|
self._queue = asyncio.Queue(maxsize=self.max_queue)
|
|
self._task = loop.create_task(self._run(self._queue))
|
|
|
|
try:
|
|
self._queue.put_nowait(record)
|
|
except asyncio.QueueFull:
|
|
self._drop("queue full")
|
|
return False
|
|
return True
|
|
|
|
def _drop(self, reason: str) -> None:
|
|
self.dropped += 1
|
|
if self.dropped == 1 or self.dropped % 1000 == 0:
|
|
logger.warning(
|
|
"Dropped tool-call record (%s); %d dropped so far", reason, self.dropped
|
|
)
|
|
|
|
async def _run(self, queue: asyncio.Queue) -> None:
|
|
while True:
|
|
batch = [await queue.get()]
|
|
while len(batch) < self.batch_size:
|
|
try:
|
|
batch.append(queue.get_nowait())
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
try:
|
|
await self._deliver(batch)
|
|
finally:
|
|
for _ in batch:
|
|
queue.task_done()
|
|
|
|
async def _deliver(self, batch: list[ToolCallRecord]) -> None:
|
|
for sink in self.sinks:
|
|
try:
|
|
await asyncio.wait_for(sink.write(batch), self.sink_timeout)
|
|
except Exception:
|
|
self.sink_errors += 1
|
|
logger.warning(
|
|
"Sink %s failed to write %d records",
|
|
type(sink).__name__,
|
|
len(batch),
|
|
exc_info=True,
|
|
)
|
|
|
|
async def flush(self) -> None:
|
|
"""Wait until every queued record has been handed to the sinks."""
|
|
if self._queue is not None and self._loop is asyncio.get_running_loop():
|
|
await self._queue.join()
|
|
|
|
async def aclose(self) -> None:
|
|
"""Flush, stop the background task and close the sinks."""
|
|
await self.flush()
|
|
if self._task is not None:
|
|
self._task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError, RuntimeError):
|
|
await self._task
|
|
self._task = None
|
|
self._queue = None
|
|
for sink in self.sinks:
|
|
close = getattr(sink, "aclose", None)
|
|
if close is not None:
|
|
try:
|
|
await close()
|
|
except Exception:
|
|
logger.warning(
|
|
"Sink %s failed to close", type(sink).__name__, exc_info=True
|
|
)
|