diff --git a/src/fabledassistant/mcp/_context.py b/src/fabledassistant/mcp/_context.py new file mode 100644 index 0000000..cd5760b --- /dev/null +++ b/src/fabledassistant/mcp/_context.py @@ -0,0 +1,20 @@ +"""Per-request MCP context. + +The ASGI middleware (mcp/server.py) populates `_user_id_ctx` from +scope['fable_user_id'] before dispatching to FastMCP tool handlers. +Tool functions then call `current_user_id()` to retrieve it. + +Tools must never fall back to a default user — `current_user_id()` raises +if called outside a properly-authed MCP request. +""" +from contextvars import ContextVar + +_user_id_ctx: ContextVar[int | None] = ContextVar("fable_mcp_user_id", default=None) + + +def current_user_id() -> int: + """Return the user_id for the in-flight MCP request. Raises if unset.""" + uid = _user_id_ctx.get() + if uid is None: + raise RuntimeError("no MCP user context — request was not bearer-authed") + return uid diff --git a/src/fabledassistant/mcp/server.py b/src/fabledassistant/mcp/server.py index c97eaac..24184b6 100644 --- a/src/fabledassistant/mcp/server.py +++ b/src/fabledassistant/mcp/server.py @@ -63,7 +63,12 @@ def mount_mcp(app: Quart) -> None: }) return scope["fable_user_id"] = user_id - await mcp_asgi(scope, receive, send) + from fabledassistant.mcp._context import _user_id_ctx + token = _user_id_ctx.set(user_id) + try: + await mcp_asgi(scope, receive, send) + finally: + _user_id_ctx.reset(token) original_asgi = app.asgi_app diff --git a/tests/test_mcp_context.py b/tests/test_mcp_context.py new file mode 100644 index 0000000..2a73903 --- /dev/null +++ b/tests/test_mcp_context.py @@ -0,0 +1,25 @@ +"""Tests for MCP per-request user_id context.""" +import pytest + +from fabledassistant.mcp._context import current_user_id, _user_id_ctx + + +def test_current_user_id_raises_when_unset(): + """No active context — must raise so tools don't silently fall back to a default.""" + _user_id_ctx.set(None) + with pytest.raises(RuntimeError, match="no MCP user context"): + current_user_id() + + +def test_current_user_id_returns_set_value(): + _user_id_ctx.set(42) + assert current_user_id() == 42 + + +def test_current_user_id_isolation_via_contextvar_token(): + """Reset via the token returned by set() restores prior state.""" + _user_id_ctx.set(1) + token = _user_id_ctx.set(2) + assert current_user_id() == 2 + _user_id_ctx.reset(token) + assert current_user_id() == 1