rsp2k
6987d947d3
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.
286 lines
9.4 KiB
Python
286 lines
9.4 KiB
Python
"""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
|