ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-24 06:24:18 +02:00
Code Issues Projects Releases Packages Wiki Activity

Scope tasks to authorization context, not session (#3800)

Browse source 
* Scope tasks to authorization context, not session

Tasks were keyed by the transport-layer Mcp-Session-Id, which is
server-assigned and changes on reconnect — so clients lost access to
their running tasks after any connection interruption.

The MCP spec says tasks should be bound to authorization context, not
session.  This replaces session_id with task_scope (derived from
AccessToken.client_id, URL-encoded) in all task data Redis keys and
Docket task keys.  When no auth is configured, a "_" sentinel is used
and security comes from UUID task ID entropy per the spec.

Session ID is still used for transport-level concerns (notification
queues, subscriber registration) and is now stored in the
TaskContextSnapshot payload so background workers can still deliver
notifications.

Also extracts all the task context infrastructure (TaskContextInfo,
TaskContextSnapshot, snapshot loading, session/server registries) from
server/dependencies.py into a new server/tasks/context.py to keep the
DI module from sprawling further.

Closes #3758

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Rename _redis_key to _snapshot_redis_key

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Document in-process session registry as an optimization

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Tidy imports and docstrings

Hoist imports where safe, keep subscriptions/notifications deferred in
handlers.py since they pull in docket at module level. Sharpen docstrings
on keys.py and context.py so each module owns its lane. Clean up the
re-export block in dependencies.py.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix misleading comment on re-export block

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Tighten task scope: include sub claim, partition keyspaces

Addresses review feedback on #3800:

- Compose task scope from client_id and the JWT sub claim (when present)
  so fixed-OAuth deployments isolate per user, not just per client.
- Replace the "_" anonymous sentinel with a tagged keyspace partition.
  Docket keys are now auth:{enc_scope}:... or anon:..., and Redis keys
  use fastmcp:task:auth:{enc_scope}:... or fastmcp:task:anon:...,
  routed through a single task_redis_prefix() helper.
- get_task_scope() returns the raw scope (or None); encoding happens
  once at the keys.py boundary, collapsing the previous double-quote
  invariant.
- Drop the dormant fallback in notifications.py that routed
  input_required relays into the anon keyspace when task_scope was
  missing -- log and skip instead.
- Add comprehensive parser/encoder tests in test_task_keys.py covering
  round-trips, malformed keys, and adversarial scopes ("anon", "_", and
  scopes containing : / | %).
- Add cross-scope rejection tests: distinct client_ids, distinct sub
  claims under a shared client_id, and authenticated vs anonymous.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Guidry 2026-04-13 13:50:43 -04:00 • committed by GitHub
commit 9a063963f4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 954 additions and 511 deletions
Download patch file Download diff file Expand all files Collapse all files

33
tests/server/tasks/test_context_background_task.py
View file

@ -23,18 +23,22 @@ from fastmcp.client.elicitation import ElicitResult
from fastmcp.dependencies import CurrentDocket
from fastmcp.server.auth import AccessToken
from fastmcp.server.context import Context
from fastmcp.server.dependencies import (
TaskContextInfo,
TaskContextSnapshot,
_set_cached_snapshot,
get_access_token,
)
from fastmcp.server.dependencies import get_access_token
from fastmcp.server.elicitation import (
AcceptedElicitation,
CancelledElicitation,
DeclinedElicitation,
)
from fastmcp.server.tasks.context import (
TaskContextInfo,
TaskContextSnapshot,
_set_cached_snapshot,
get_task_scope,
)
from fastmcp.server.tasks.elicitation import handle_task_input
from fastmcp.server.tasks.keys import (
task_redis_prefix,
)
# =============================================================================
# Unit tests: Context API surface (no Redis/Docket needed)
@ -262,7 +266,8 @@ class TestBackgroundTaskIntegration:
assert origin != ""
# Verify the snapshot in Redis contains the same value
key = docket.key(f"fastmcp:task:{ctx.session_id}:{ctx.task_id}:snapshot")
task_scope = get_task_scope()
key = docket.key(f"{task_redis_prefix(task_scope)}:{ctx.task_id}:snapshot")
async with docket.redis() as redis:
raw = await redis.get(key)
@ -361,7 +366,7 @@ class TestBackgroundTaskIntegration:
# Task already completed — no elicitation waiting
success = await handle_task_input(
task_id=task.task_id,
session_id="nonexistent-session",
task_scope="nonexistent-scope",
action="accept",
content={"value": "too late"},
fastmcp=mcp,
@ -429,9 +434,9 @@ class TestAccessTokenInBackgroundTasks:
"test-task",
TaskContextSnapshot(access_token_json=expired.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
):
assert get_access_token() is None
@ -447,9 +452,9 @@ class TestAccessTokenInBackgroundTasks:
"test-task",
TaskContextSnapshot(access_token_json=valid.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
):
result = get_access_token()
assert result is not None
@ -466,9 +471,9 @@ class TestAccessTokenInBackgroundTasks:
"test-task",
TaskContextSnapshot(access_token_json=no_expiry.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
):
result = get_access_token()
assert result is not None

170
tests/server/tasks/test_task_keys.py Normal file
View file

@ -0,0 +1,170 @@
"""Tests for ``fastmcp.server.tasks.keys`` — the encoding boundary that
separates authenticated and anonymous task keyspaces.
Cross-scope isolation depends on these encodings being unambiguous and
round-trippable, so the tests cover: tag dispatch (``auth``/``anon``),
the ``None`` ⇄ anonymous round trip, encoding of values that contain the
``:`` delimiter, error paths for malformed keys, and the parity between
the Docket-key prefix and the Redis-key prefix.
"""
import pytest
from fastmcp.server.tasks.keys import (
build_task_key,
get_client_task_id_from_key,
parse_task_key,
task_redis_prefix,
)
ROUND_TRIP_CASES = [
("client-a", "task-1", "tool", "my_tool"),
(None, "task-1", "tool", "my_tool"),
("client-a", "task-1", "resource", "file://data.txt"),
(None, "task-1", "resource", "file://data.txt"),
("client-a", "task-1", "template", "users://{id}"),
("client-a", "task-1", "prompt", "greet@1.0.0"),
# Scope contains the inner separator used by get_task_scope (client_id|sub).
("client|sub-42", "task-1", "tool", "my_tool"),
# Adversarial: scope is literally the anon tag — must not collide.
("anon", "task-1", "tool", "my_tool"),
# Adversarial: scope is literally the legacy "_" sentinel.
("_", "task-1", "tool", "my_tool"),
# Scope contains every delimiter we care about.
("a:b/c d%e|f", "task-1", "tool", "my_tool"),
# Component identifier with colons, slashes, percent, spaces.
("client-a", "task-1", "resource", "https://x/y?z=1&q=a b"),
# UUID-shaped task id (the realistic case).
("client-a", "0c3e9b14-3a3f-4b3a-9b1a-1d8d6e6e0c11", "tool", "t"),
]
@pytest.mark.parametrize(
("scope", "task_id", "task_type", "identifier"), ROUND_TRIP_CASES
)
def test_round_trip_preserves_all_fields(
scope: str | None, task_id: str, task_type: str, identifier: str
):
key = build_task_key(scope, task_id, task_type, identifier)
parsed = parse_task_key(key)
assert parsed == {
"task_scope": scope,
"client_task_id": task_id,
"task_type": task_type,
"component_identifier": identifier,
}
@pytest.mark.parametrize(
("scope", "task_id", "task_type", "identifier"), ROUND_TRIP_CASES
)
def test_get_client_task_id_round_trip(
scope: str | None, task_id: str, task_type: str, identifier: str
):
key = build_task_key(scope, task_id, task_type, identifier)
assert get_client_task_id_from_key(key) == task_id
def test_authenticated_key_uses_auth_tag():
key = build_task_key("client-a", "task-1", "tool", "my_tool")
assert key.startswith("auth:")
assert key == "auth:client-a:task-1:tool:my_tool"
def test_anonymous_key_uses_anon_tag():
key = build_task_key(None, "task-1", "tool", "my_tool")
assert key.startswith("anon:")
assert key == "anon:task-1:tool:my_tool"
def test_anonymous_and_literal_anon_scope_have_disjoint_keyspaces():
"""A real anonymous task and a (hostile) authenticated task whose scope
literally equals "anon" must not collide."""
anon_key = build_task_key(None, "task-1", "tool", "x")
impostor_key = build_task_key("anon", "task-1", "tool", "x")
assert anon_key != impostor_key
assert parse_task_key(anon_key)["task_scope"] is None
assert parse_task_key(impostor_key)["task_scope"] == "anon"
def test_legacy_underscore_scope_is_just_a_string_now():
"""Belt-and-suspenders: a client_id of "_" no longer aliases anonymous."""
underscore_key = build_task_key("_", "task-1", "tool", "x")
anon_key = build_task_key(None, "task-1", "tool", "x")
assert underscore_key != anon_key
assert parse_task_key(underscore_key)["task_scope"] == "_"
def test_component_identifier_with_colons_is_recovered():
key = build_task_key("client-a", "task-1", "resource", "file://data:special.txt")
assert parse_task_key(key)["component_identifier"] == "file://data:special.txt"
def test_scope_with_colons_is_recovered():
key = build_task_key("a:b:c", "task-1", "tool", "t")
parsed = parse_task_key(key)
assert parsed["task_scope"] == "a:b:c"
assert parsed["client_task_id"] == "task-1"
def test_scope_pipe_separator_is_preserved():
"""``get_task_scope`` composes ``client_id|sub`` — the ``|`` must survive."""
key = build_task_key("client-a|user-42", "task-1", "tool", "t")
assert parse_task_key(key)["task_scope"] == "client-a|user-42"
@pytest.mark.parametrize(
"bad_key",
[
"",
"client-a:task-1:tool:my_tool", # legacy untagged format
"weird:client-a:task-1:tool:my_tool", # unknown tag
"auth:client-a:task-1:tool", # missing identifier
"auth:client-a", # truncated
"anon:task-1:tool", # truncated anon
"anon", # tag only
"auth", # tag only
":task-1:tool:t", # empty tag
],
)
def test_parse_rejects_malformed_keys(bad_key: str):
with pytest.raises(ValueError):
parse_task_key(bad_key)
def test_redis_prefix_authenticated():
assert task_redis_prefix("client-a") == "fastmcp:task:auth:client-a"
def test_redis_prefix_anonymous():
assert task_redis_prefix(None) == "fastmcp:task:anon"
def test_redis_prefix_disjoint_for_anon_vs_literal_anon_scope():
assert task_redis_prefix(None) != task_redis_prefix("anon")
def test_redis_prefix_disjoint_for_anon_vs_literal_underscore_scope():
assert task_redis_prefix(None) != task_redis_prefix("_")
def test_redis_prefix_encodes_special_characters():
# Colons, slashes, pipes in the scope must not break the prefix shape.
prefix = task_redis_prefix("client:a/b|sub")
assert prefix.startswith("fastmcp:task:auth:")
# Exactly four ":" delimiters: fastmcp / task / auth / encoded-scope.
assert prefix.count(":") == 3
def test_docket_and_redis_prefixes_agree_on_partition():
"""The Docket key tag and the Redis prefix tag must always match — that is
the load-bearing invariant for cross-scope isolation."""
auth_docket = build_task_key("client-a", "task-1", "tool", "x")
auth_redis = task_redis_prefix("client-a")
assert auth_docket.split(":", 1)[0] == "auth"
assert ":auth:" in auth_redis
anon_docket = build_task_key(None, "task-1", "tool", "x")
anon_redis = task_redis_prefix(None)
assert anon_docket.split(":", 1)[0] == "anon"
assert anon_redis.endswith(":anon")

144
tests/server/tasks/test_task_security.py
View file

@ -1,22 +1,27 @@
"""
Tests for session-based task ID isolation (CRITICAL SECURITY).
Tests for authorization-based task isolation (CRITICAL SECURITY).
Ensures that tasks are properly scoped to sessions and clients cannot
access each other's tasks.
Ensures that tasks are properly scoped to authorization identity and clients
cannot access each other's tasks.
"""
import pytest
from mcp.server.auth.middleware.auth_context import (
auth_context_var,
)
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.server.auth import AccessToken
@pytest.fixture
async def task_server():
def task_server():
"""Create a server with background tasks enabled."""
mcp = FastMCP("security-test-server")
@mcp.tool(task=True) # Enable background execution
@mcp.tool(task=True)
async def secret_tool(data: str) -> str:
"""A tool that processes sensitive data."""
return f"Secret result: {data}"
@ -24,24 +29,121 @@ async def task_server():
return mcp
async def test_same_session_can_access_all_its_tasks(task_server):
"""A single session can access all tasks it created."""
async def test_same_client_can_access_all_its_tasks(task_server: FastMCP):
"""A single authenticated client can access all tasks it created."""
token = AccessToken(
token="token-a",
client_id="client-a",
scopes=["read"],
)
reset = auth_context_var.set(AuthenticatedUser(token))
try:
async with Client(task_server) as client:
task1 = await client.call_tool(
"secret_tool", {"data": "first"}, task=True, task_id="task-1"
)
task2 = await client.call_tool(
"secret_tool", {"data": "second"}, task=True, task_id="task-2"
)
await task1.wait(timeout=2.0)
await task2.wait(timeout=2.0)
result1 = await task1.result()
result2 = await task2.result()
assert "first" in str(result1.data)
assert "second" in str(result2.data)
finally:
auth_context_var.reset(reset)
async def test_unauthenticated_client_can_access_its_tasks(task_server: FastMCP):
"""An unauthenticated client can access tasks it created (by task ID)."""
async with Client(task_server) as client:
# Submit multiple tasks
task1 = await client.call_tool(
"secret_tool", {"data": "first"}, task=True, task_id="task-1"
)
task2 = await client.call_tool(
"secret_tool", {"data": "second"}, task=True, task_id="task-2"
task = await client.call_tool(
"secret_tool", {"data": "hello"}, task=True, task_id="my-task"
)
await task.wait(timeout=2.0)
result = await task.result()
assert "hello" in str(result.data)
# Wait for both to complete
await task1.wait(timeout=2.0)
await task2.wait(timeout=2.0)
# Should be able to access both
result1 = await task1.result()
result2 = await task2.result()
def _set_auth(client_id: str, sub: str | None = None):
"""Install an auth context for a given client_id/sub. Returns the reset token."""
claims = {"sub": sub} if sub else {}
token = AccessToken(
token=f"token-{client_id}-{sub or ''}",
client_id=client_id,
scopes=["read"],
claims=claims,
)
return auth_context_var.set(AuthenticatedUser(token))
assert "first" in str(result1.data)
assert "second" in str(result2.data)
async def _submit_task_id(client: Client, data: str) -> str:
"""Submit a background task and return its server-assigned task id."""
task = await client.call_tool("secret_tool", {"data": data}, task=True)
await task.wait(timeout=2.0)
return task.task_id
async def test_distinct_clients_cannot_access_each_others_tasks(
task_server: FastMCP,
):
"""Two distinct authenticated clients live in disjoint scopes — looking up
a peer's task id returns 'not found'."""
reset = _set_auth("client-a")
try:
async with Client(task_server) as client_a:
task_id = await _submit_task_id(client_a, "client-a-secret")
finally:
auth_context_var.reset(reset)
reset = _set_auth("client-b")
try:
async with Client(task_server) as client_b:
with pytest.raises(Exception, match="not found"):
await client_b.get_task_status(task_id)
finally:
auth_context_var.reset(reset)
async def test_distinct_subs_same_client_id_cannot_access_each_others_tasks(
task_server: FastMCP,
):
"""Fixed-OAuth case: two users share a client_id but have distinct ``sub``
claims. The ``sub``-aware scope must still isolate them."""
shared_client = "shared-oauth-app"
reset = _set_auth(shared_client, sub="user-alice")
try:
async with Client(task_server) as alice:
task_id = await _submit_task_id(alice, "alice-secret")
finally:
auth_context_var.reset(reset)
reset = _set_auth(shared_client, sub="user-bob")
try:
async with Client(task_server) as bob:
with pytest.raises(Exception, match="not found"):
await bob.get_task_status(task_id)
finally:
auth_context_var.reset(reset)
async def test_authenticated_and_anonymous_keyspaces_are_disjoint(
task_server: FastMCP,
):
"""An anonymous client must not be able to read an authenticated client's
tasks (and vice versa) even when colliding on task id."""
reset = _set_auth("client-a")
try:
async with Client(task_server) as authed:
authed_task_id = await _submit_task_id(authed, "authed-secret")
finally:
auth_context_var.reset(reset)
async with Client(task_server) as anon:
with pytest.raises(Exception, match="not found"):
await anon.get_task_status(authed_task_id)