From 6987d947d31432f771a94ec031819455bcfc3888 Mon Sep 17 00:00:00 2001 From: Ryan Malloy Date: Thu, 1 Oct 2026 23:57:07 -0600 Subject: [PATCH] Run the identity resolver again after the tool when its first answer is incomplete Since 2026.10.01.4 the resolver runs before the tool, so a tool that creates the caller's identity (blender-mcp's register_client, which makes the bus client the resolver looks up) recorded its own row without it. When the first answer lacks user_sub, caller_kind or client_id, or any of the new late_identity_keys, the resolver runs once more after the tool, alongside the enricher, for calls that are kept. Its values fill the gaps; the first answer wins on conflict, matching what events during the call saw. --- Makefile | 2 +- .../instrumentation/middleware.py | 93 ++++-- tests/unit/test_consistency.py | 3 +- tests/unit/test_late_identity.py | 285 ++++++++++++++++++ 4 files changed, 363 insertions(+), 20 deletions(-) create mode 100644 tests/unit/test_late_identity.py diff --git a/Makefile b/Makefile index a8e086e..51f8408 100644 --- a/Makefile +++ b/Makefile @@ -61,4 +61,4 @@ lint: uv run --no-sync ruff check src/ examples/ tests/_compat.py tests/_databases.py \ tests/unit/test_test_databases.py tests/unit/test_instrumentation.py \ tests/unit/test_embeddings.py tests/unit/test_backfill.py \ - tests/unit/test_consistency.py + tests/unit/test_consistency.py tests/unit/test_late_identity.py diff --git a/src/fastmcp_feedback/instrumentation/middleware.py b/src/fastmcp_feedback/instrumentation/middleware.py index 683911f..130c39b 100644 --- a/src/fastmcp_feedback/instrumentation/middleware.py +++ b/src/fastmcp_feedback/instrumentation/middleware.py @@ -140,11 +140,17 @@ class InstrumentationMiddleware(Middleware): meta_only_tools: tools that are never recorded above ``meta``. identity_resolver: ``(context) -> dict`` returning ``user_sub``, ``caller_kind`` and ``client_id`` (stored as columns) plus any - other keys (stored in ``identity``). May be async. Runs once per - recorded call, before the tool, so events recorded during the - call get the same identity. It also runs for calls that sampling - later drops, and an async one can delay the tool by up to - ``hook_timeout``; keep it cheap. + other keys (stored in ``identity``). May be async. Runs before + the tool, so events recorded during the call get that identity; + it also runs for calls that sampling later drops, and an async + one can delay the tool by up to ``hook_timeout``; keep it cheap. + When that first answer is incomplete (any of the three columns + missing, or any of ``late_identity_keys``), it runs again after + the tool, alongside the enricher, and the call's own record + takes the second answer's values where the first had none. + late_identity_keys: identity keys that may only exist once the tool + has run, such as an id a registration or login tool creates. A + first answer without one of them gets the second run. enricher: ``(tool, args, result, context) -> dict`` of extra fields, stored in ``extra``. A ``server_version`` key fills that column. May be async. @@ -189,6 +195,7 @@ class InstrumentationMiddleware(Middleware): redactor: Redactor | None = None, meta_only_tools: Iterable[str] = (), identity_resolver: IdentityResolver | None = None, + late_identity_keys: Iterable[str] = (), enricher: Enricher | None = None, server_version: str | None = None, hook_timeout: float = 0.25, @@ -213,6 +220,7 @@ class InstrumentationMiddleware(Middleware): self.redactor = redactor or Redactor() self.meta_only_tools = frozenset(meta_only_tools) self.identity_resolver = identity_resolver + self.late_identity_keys = tuple(late_identity_keys) self.enricher = enricher self.server_version = server_version self.hook_timeout = hook_timeout @@ -273,9 +281,9 @@ class InstrumentationMiddleware(Middleware): scope = CallScope( id=str(uuid.uuid4()), session_id=session_key(context.fastmcp_context), - user_sub=who.columns.get("user_sub"), - caller_kind=who.columns.get("caller_kind"), - client_id=who.columns.get("client_id"), + user_sub=who.column("user_sub"), + caller_kind=who.column("caller_kind"), + client_id=who.column("client_id"), ) token = self._call_scope.set(scope) except Exception: @@ -336,21 +344,30 @@ class InstrumentationMiddleware(Middleware): self.dispatcher.submit(record) async def _resolve_identity(self, context: MiddlewareContext) -> "_Identity": - """Run the identity resolver once for a call. Never raises. + """Run the identity resolver once. Never raises. Returns the redacted result split into the three columns and the - rest. ``client_id`` falls back to the FastMCP context's. + rest, plus the FastMCP context's ``client_id`` as the fallback for a + resolver that gives none. """ + fallback = None try: + fallback = _attr(context.fastmcp_context, "client_id") identity = await self._call_hook(self.identity_resolver, context) identity = self.redactor.redact(identity) columns = {k: identity.pop(k) for k in _IDENTITY_COLUMNS if k in identity} - if not columns.get("client_id"): - columns["client_id"] = _attr(context.fastmcp_context, "client_id") - return _Identity(columns, identity) + return _Identity(columns, identity, fallback) except Exception: logger.warning("Failed to resolve identity", exc_info=True) - return _Identity({}, {}) + return _Identity({}, {}, fallback) + + def _identity_incomplete(self, who: "_Identity") -> bool: + """Whether the resolver's first answer lacks something it may know later.""" + if self.identity_resolver is None: + return False + return any(not who.column(k) for k in _IDENTITY_COLUMNS) or any( + who.rest.get(k) is None for k in self.late_identity_keys + ) async def _call_hook(self, hook: Callable | None, *args: Any) -> dict[str, Any]: if hook is None: @@ -395,9 +412,17 @@ class InstrumentationMiddleware(Middleware): args = getattr(context.message, "arguments", None) or {} ctx = context.fastmcp_context - columns = who.columns + enrich = self._call_hook(self.enricher, tool, args, payload, context) + if self._identity_incomplete(who): + # A tool can create the identity it runs under (a registration, + # a login). Ask again, concurrently with the enricher so the + # record is not built any later than the slower of the two. + later, extra = await asyncio.gather(self._resolve_identity(context), enrich) + who = who.filled_from(later) + else: + extra = await enrich identity = dict(who.rest) - extra = await self._call_hook(self.enricher, tool, args, payload, context) + columns = {k: who.column(k) for k in _IDENTITY_COLUMNS} extra = self.redactor.redact(extra) server_version = extra.pop("server_version", None) or self.server_version @@ -828,10 +853,42 @@ class InstrumentationMiddleware(Middleware): class _Identity(NamedTuple): - """One call's resolved identity: the column values and everything else.""" + """One answer from the identity resolver, redacted. - columns: dict[str, str | None] + ``columns`` holds whichever of user_sub, caller_kind and client_id it + returned, ``rest`` everything else, and ``fallback_client_id`` the + FastMCP context's client_id for when it returned none. + """ + + columns: dict[str, Any] rest: dict[str, Any] + fallback_client_id: str | None = None + + def column(self, key: str) -> Any: + value = self.columns.get(key) + if key == "client_id" and not value: + return self.fallback_client_id + return value + + def filled_from(self, later: "_Identity") -> "_Identity": + """This answer, with gaps and None values filled from ``later``. + + Values already here win, so the call's record agrees with the events + recorded during the call, which saw only this answer. + """ + + def fill(first: dict[str, Any], second: dict[str, Any]) -> dict[str, Any]: + merged = dict(first) + for key, value in second.items(): + if merged.get(key) is None: + merged[key] = value + return merged + + return _Identity( + fill(self.columns, later.columns), + fill(self.rest, later.rest), + self.fallback_client_id or later.fallback_client_id, + ) def _is_sink(sink: Any, module: str, name: str) -> bool: diff --git a/tests/unit/test_consistency.py b/tests/unit/test_consistency.py index 5e9140e..86dea4f 100644 --- a/tests/unit/test_consistency.py +++ b/tests/unit/test_consistency.py @@ -364,7 +364,8 @@ class TestEventIdentity: [ev] = sink.events assert rec.ok and ev.call_id == rec.id assert ev.user_sub is None and rec.user_sub is None - assert calls == [1] + # An empty first answer gets a second run after the tool; it fails too. + assert calls == [1, 1] async def test_slow_resolver_cut_off_before_tool(self): async def slow(context): diff --git a/tests/unit/test_late_identity.py b/tests/unit/test_late_identity.py new file mode 100644 index 0000000..352e2bb --- /dev/null +++ b/tests/unit/test_late_identity.py @@ -0,0 +1,285 @@ +"""The identity resolver's second run, after the tool, for identity the tool creates. + +Modelled on blender-mcp: an addon calls register_client, and only after that +call can the resolver find the bus client (bus_client_uuid, bus_role) for the +caller's session. The first run, before the tool, already knows the user. +""" + +import asyncio +import contextlib +import time + +from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError + +from fastmcp_feedback.instrumentation import REDACTED, MemorySink, instrument + + +class Bus: + """Who is registered, as the tool leaves it; read by the resolvers.""" + + def __init__(self): + self.client = None + + +class BlenderResolver: + """blender-mcp's identify(): user and role always, bus keys once registered.""" + + def __init__(self, bus): + self.bus = bus + self.calls = 0 + + def __call__(self, context): + self.calls += 1 + out = {"user_sub": "user-1", "caller_kind": "addon", "client_id": "dcr-9"} + if self.bus.client is not None: + out["bus_client_uuid"] = self.bus.client + out["bus_role"] = "interactive" + return out + + +class EmptyUntilRegistered: + """Nothing before registration, every field after.""" + + def __init__(self, bus): + self.bus = bus + self.calls = 0 + + def __call__(self, context): + self.calls += 1 + if self.bus.client is None: + return {} + return { + "user_sub": "user-1", + "caller_kind": "addon", + "client_id": "dcr-9", + "bus_client_uuid": self.bus.client, + } + + +def make_app(resolver_cls, *, late=(), sample=None, **kwargs): + bus = Bus() + resolver = resolver_cls(bus) + sink = MemorySink() + app = FastMCP("late-identity") + holder = {} + + @app.tool + async def register_client(client_uuid: str) -> dict: + holder["mw"].record_event("bus.registering", key=client_uuid) + bus.client = client_uuid + return {"status": "ok"} + + @app.tool + async def register_then_fail(client_uuid: str) -> dict: + bus.client = client_uuid + raise RuntimeError("registered, then broke") + + @app.tool + async def quiet() -> str: + return "ok" + + @app.tool + async def ping() -> str: + holder["mw"].record_event("bus.ping") + return "pong" + + kwargs.setdefault("identity_resolver", resolver) + if sample is not None: + kwargs.update( + sample_rates={"ping": sample, "quiet": sample}, random=lambda: 0.99 + ) + mw = instrument(app, [sink], late_identity_keys=late, **kwargs) + holder["mw"] = mw + return app, sink, mw, resolver + + +async def call(app, mw, *calls): + async with Client(app) as client: + for tool, args in calls: + with contextlib.suppress(ToolError): + await client.call_tool(tool, args) + await mw.flush() + + +def columns(record): + return (record.user_sub, record.caller_kind, record.client_id) + + +class TestRegistration: + async def test_blender_shaped_registration_row_gets_bus_identity(self): + app, sink, mw, who = make_app(BlenderResolver, late=("bus_client_uuid",)) + await call(app, mw, ("register_client", {"client_uuid": "c-1"})) + [rec] = sink.calls + assert columns(rec) == ("user-1", "addon", "dcr-9") + assert rec.identity == {"bus_client_uuid": "c-1", "bus_role": "interactive"} + assert who.calls == 2 + # The event recorded during the call saw the first answer only. + [ev] = sink.events + assert ev.call_id == rec.id and columns(ev) == ("user-1", "addon", "dcr-9") + + async def test_blender_shaped_without_late_keys_runs_once(self): + # The three columns are all there before the tool, so nothing marks + # the answer incomplete: the bus keys need late_identity_keys. + app, sink, mw, who = make_app(BlenderResolver) + await call(app, mw, ("register_client", {"client_uuid": "c-1"})) + [rec] = sink.calls + assert rec.identity == {} and who.calls == 1 + + async def test_registered_caller_runs_once(self): + app, sink, mw, who = make_app(BlenderResolver, late=("bus_client_uuid",)) + await call( + app, + mw, + ("register_client", {"client_uuid": "c-1"}), + ("ping", {}), + ("ping", {}), + ) + assert who.calls == 2 + 1 + 1 + assert [r.identity.get("bus_client_uuid") for r in sink.calls] == [ + "c-1", + "c-1", + "c-1", + ] + + async def test_empty_first_answer_filled_after(self): + app, sink, mw, who = make_app(EmptyUntilRegistered) + await call(app, mw, ("register_client", {"client_uuid": "c-2"})) + [rec] = sink.calls + [ev] = sink.events + assert columns(rec) == ("user-1", "addon", "dcr-9") + assert rec.identity == {"bus_client_uuid": "c-2"} + assert columns(ev) == (None, None, None) # not rewritten + assert who.calls == 2 + + async def test_tool_that_raises_gets_second_run_too(self): + app, sink, mw, who = make_app(EmptyUntilRegistered) + await call(app, mw, ("register_then_fail", {"client_uuid": "c-3"})) + [rec] = sink.calls + assert rec.outcome == "error" + assert columns(rec) == ("user-1", "addon", "dcr-9") + assert who.calls == 2 + + +class TestMerge: + async def test_first_answer_wins_on_conflict(self): + answers = iter( + [ + {"user_sub": "early", "caller_kind": None, "org": None}, + { + "user_sub": "late", + "caller_kind": "k", + "client_id": "c", + "org": "acme", + "team": "qa", + }, + ] + ) + app, sink, mw, _ = make_app( + EmptyUntilRegistered, identity_resolver=lambda c: next(answers) + ) + await call(app, mw, ("ping", {})) + [rec] = sink.calls + [ev] = sink.events + assert columns(rec) == ("early", "k", "c") + assert rec.identity == {"org": "acme", "team": "qa"} + assert columns(ev) == ("early", None, None) + + async def test_second_answer_redacted(self): + token = "bmcp_abcdefgh12345678" + answers = iter([{}, {"user_sub": "u", "caller_kind": f"pat:{token}"}]) + app, sink, mw, _ = make_app( + EmptyUntilRegistered, identity_resolver=lambda c: next(answers) + ) + await call(app, mw, ("ping", {})) + assert sink.calls[0].caller_kind == f"pat:{REDACTED}" + + +class TestFailures: + async def test_second_run_raising_keeps_first_answer(self): + runs = [] + + def who(context): + runs.append(1) + if len(runs) == 1: + return {"user_sub": "u1"} + raise RuntimeError("resolver bug") + + app, sink, mw, _ = make_app(EmptyUntilRegistered, identity_resolver=who) + await call(app, mw, ("ping", {})) + [rec] = sink.calls + assert rec.ok and rec.user_sub == "u1" and len(runs) == 2 + + async def test_second_run_timing_out_keeps_first_answer(self): + runs = [] + + async def who(context): + runs.append(1) + if len(runs) == 1: + return {"user_sub": "u1"} + await asyncio.sleep(5) + return {"user_sub": "late", "caller_kind": "k"} + + app, sink, mw, _ = make_app( + EmptyUntilRegistered, identity_resolver=who, hook_timeout=0.05 + ) + await call(app, mw, ("ping", {})) + [rec] = sink.calls + assert rec.ok and rec.user_sub == "u1" and rec.caller_kind is None + + +class TestWhenItRuns: + async def test_complete_answer_runs_once(self): + app, sink, mw, who = make_app(BlenderResolver) + await call(app, mw, ("ping", {}), ("ping", {})) + assert who.calls == 2 and len(sink.calls) == 2 + + async def test_sampled_out_call_makes_no_second_run(self): + app, sink, mw, who = make_app(EmptyUntilRegistered, sample=0.1) + await call(app, mw, ("quiet", {})) + assert sink.calls == [] and who.calls == 1 + + async def test_sampled_call_kept_for_its_event_gets_second_run(self): + app, sink, mw, who = make_app(EmptyUntilRegistered, sample=0.1) + await call(app, mw, ("ping", {})) + assert len(sink.calls) == 1 and who.calls == 2 + + async def test_no_resolver_no_runs(self): + app, sink, mw, _ = make_app(EmptyUntilRegistered, identity_resolver=None) + await call(app, mw, ("ping", {})) + [rec] = sink.calls + assert columns(rec) == (None, None, None) + + async def test_second_run_overlaps_enricher(self): + # The record is built before the result reaches the client, so the + # second run must not add its time to the enricher's. + runs = [] + + async def who(context): + runs.append(1) + if len(runs) > 1: + await asyncio.sleep(0.4) + return {} + + async def enrich(tool, args, result, context): + await asyncio.sleep(0.4) + return {} + + app, sink, mw, _ = make_app( + EmptyUntilRegistered, + identity_resolver=who, + enricher=enrich, + hook_timeout=2, + ) + async with Client(app) as client: + start = time.perf_counter() + await client.call_tool("ping", {}) + elapsed = time.perf_counter() - start + assert len(runs) == 2 + assert elapsed < 0.75, elapsed + + +async def test_mode_off_runs_nothing(): + app, sink, mw, who = make_app(EmptyUntilRegistered, mode="off") + await call(app, mw, ("ping", {})) + assert who.calls == 0