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

@@ -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

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:

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):

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