60e0ffc959
Upgrade checks / Notify on failure (push) Has been cancelled
Upgrade checks / Close issue on success (push) Has been cancelled
Schema Crash Test / Real-world schema crash test (232K schemas) (push) Has been cancelled
Run static analysis / static_analysis (push) Has been cancelled
Tests / Tests: Python 3.10 on ubuntu-latest (push) Has been cancelled
Tests / Tests: Python 3.13 on ubuntu-latest (push) Has been cancelled
Tests / Tests: Python 3.10 on windows-latest (push) Has been cancelled
Tests / Tests with lowest-direct dependencies (push) Has been cancelled
Tests / MCP conformance tests (push) Has been cancelled
Tests / Integration tests (push) Has been cancelled
Tests / Package install smoke (push) Has been cancelled
Upgrade checks / Static analysis (push) Has been cancelled
Upgrade checks / Tests: Python 3.10 on ubuntu-latest (push) Has been cancelled
Upgrade checks / Tests: Python 3.13 on ubuntu-latest (push) Has been cancelled
Upgrade checks / Tests: Python 3.10 on windows-latest (push) Has been cancelled
Upgrade checks / Integration tests (push) Has been cancelled
Update MCPServerConfig Schema / update-config-schema (push) Has been cancelled
Update SDK Documentation / update-sdk-docs (push) Has been cancelled
770 lines
27 KiB
Python
770 lines
27 KiB
Python
"""Tests for session-specific visibility control via Context."""
|
|
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
|
|
import anyio
|
|
import mcp_types
|
|
|
|
from fastmcp.client.messages import MessageHandler
|
|
from fastmcp.server.context import Context
|
|
from fastmcp.server.server import FastMCP
|
|
|
|
|
|
@dataclass
|
|
class NotificationRecording:
|
|
"""Record of a notification that was received."""
|
|
|
|
method: str
|
|
notification: mcp_types.ServerNotification
|
|
timestamp: datetime = field(default_factory=datetime.now)
|
|
|
|
|
|
class RecordingMessageHandler(MessageHandler):
|
|
"""A message handler that records all notifications."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.notifications: list[NotificationRecording] = []
|
|
|
|
async def on_notification(self, message: mcp_types.ServerNotification) -> None:
|
|
"""Record all notifications with timestamp."""
|
|
# SDK v2 delivers notifications unwrapped (no `.root` wrapper).
|
|
self.notifications.append(
|
|
NotificationRecording(method=message.method, notification=message)
|
|
)
|
|
|
|
def get_notifications(
|
|
self, method: str | None = None
|
|
) -> list[NotificationRecording]:
|
|
"""Get all recorded notifications, optionally filtered by method."""
|
|
if method is None:
|
|
return self.notifications
|
|
return [n for n in self.notifications if n.method == method]
|
|
|
|
def reset(self):
|
|
"""Clear all recorded notifications."""
|
|
self.notifications.clear()
|
|
|
|
|
|
class TestSessionVisibility:
|
|
"""Test session-specific visibility control via Context."""
|
|
|
|
async def test_enable_components_stores_rule_dict(self):
|
|
"""Test that enable_components stores a rule dict in session state."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
# Check that the rule was stored
|
|
rules = await ctx._get_visibility_rules()
|
|
assert len(rules) == 1
|
|
assert rules[0]["enabled"] is True
|
|
assert rules[0]["tags"] == ["finance"]
|
|
return "activated"
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("activate_finance", {})
|
|
assert result.data == "activated"
|
|
|
|
async def test_disable_components_stores_rule_dict(self):
|
|
"""Test that disable_components stores a rule dict in session state."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"internal"})
|
|
def internal_tool() -> str:
|
|
return "internal"
|
|
|
|
@mcp.tool
|
|
async def deactivate_internal(ctx: Context) -> str:
|
|
await ctx.disable_components(tags={"internal"})
|
|
# Check that the rule was stored
|
|
rules = await ctx._get_visibility_rules()
|
|
assert len(rules) == 1
|
|
assert rules[0]["enabled"] is False
|
|
assert rules[0]["tags"] == ["internal"]
|
|
return "deactivated"
|
|
|
|
async with Client(mcp) as client:
|
|
result = await client.call_tool("deactivate_internal", {})
|
|
assert result.data == "deactivated"
|
|
|
|
async def test_session_rules_override_global_disables(self):
|
|
"""Test that session enable rules override global disable transforms."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
# Globally disable finance tools
|
|
mcp.disable(tags={"finance"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Before activation, finance tool should not be visible
|
|
tools_before = await client.list_tools()
|
|
assert not any(t.name == "finance_tool" for t in tools_before)
|
|
|
|
# Activate finance for this session
|
|
await client.call_tool("activate_finance", {})
|
|
|
|
# After activation, finance tool should be visible in this session
|
|
tools_after = await client.list_tools()
|
|
assert any(t.name == "finance_tool" for t in tools_after)
|
|
|
|
async def test_rules_persist_across_requests(self):
|
|
"""Test that session rules persist across multiple requests."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
@mcp.tool
|
|
async def check_rules(ctx: Context) -> int:
|
|
rules = await ctx._get_visibility_rules()
|
|
return len(rules)
|
|
|
|
# Globally disable finance tools
|
|
mcp.disable(tags={"finance"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Activate finance
|
|
await client.call_tool("activate_finance", {})
|
|
|
|
# In a subsequent request, rules should still be there
|
|
result = await client.call_tool("check_rules", {})
|
|
assert result.data == 1
|
|
|
|
# And finance tool should still be visible
|
|
tools = await client.list_tools()
|
|
assert any(t.name == "finance_tool" for t in tools)
|
|
|
|
async def test_rules_isolated_between_sessions(self):
|
|
"""Test that session rules are isolated between different sessions."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
# Globally disable finance tools
|
|
mcp.disable(tags={"finance"})
|
|
|
|
# Session A activates finance
|
|
async with Client(mcp) as client_a:
|
|
await client_a.call_tool("activate_finance", {})
|
|
tools_a = await client_a.list_tools()
|
|
assert any(t.name == "finance_tool" for t in tools_a)
|
|
|
|
# Session B should not see finance tool (different session)
|
|
async with Client(mcp) as client_b:
|
|
tools_b = await client_b.list_tools()
|
|
assert not any(t.name == "finance_tool" for t in tools_b)
|
|
|
|
async def test_version_spec_serialization(self):
|
|
"""Test that VersionSpec is serialized/deserialized correctly."""
|
|
from fastmcp import Client
|
|
from fastmcp.utilities.versions import VersionSpec
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(version="1.0.0")
|
|
def old_tool() -> str:
|
|
return "old"
|
|
|
|
@mcp.tool(version="2.0.0")
|
|
def new_tool() -> str:
|
|
return "new"
|
|
|
|
@mcp.tool
|
|
async def enable_v2_only(ctx: Context) -> str:
|
|
await ctx.enable_components(version=VersionSpec(gte="2.0.0"))
|
|
# Check serialization - version is stored as a dict
|
|
rules = await ctx._get_visibility_rules()
|
|
assert rules[0]["version"]["gte"] == "2.0.0"
|
|
assert rules[0]["version"]["lt"] is None
|
|
assert rules[0]["version"]["eq"] is None
|
|
return "enabled"
|
|
|
|
# Globally disable all versioned tools
|
|
mcp.disable(names={"old_tool", "new_tool"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Enable v2 tools
|
|
await client.call_tool("enable_v2_only", {})
|
|
|
|
# Should see new_tool (v2.0.0) but not old_tool (v1.0.0)
|
|
tools = await client.list_tools()
|
|
assert any(t.name == "new_tool" for t in tools)
|
|
assert not any(t.name == "old_tool" for t in tools)
|
|
|
|
async def test_clear_visibility_rules(self):
|
|
"""Test that reset_visibility removes all session rules."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
@mcp.tool
|
|
async def clear_rules(ctx: Context) -> str:
|
|
await ctx.reset_visibility()
|
|
rules = await ctx._get_visibility_rules()
|
|
assert len(rules) == 0
|
|
return "cleared"
|
|
|
|
# Globally disable finance tools
|
|
mcp.disable(tags={"finance"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Activate finance
|
|
await client.call_tool("activate_finance", {})
|
|
tools_after_activate = await client.list_tools()
|
|
assert any(t.name == "finance_tool" for t in tools_after_activate)
|
|
|
|
# Clear rules
|
|
await client.call_tool("clear_rules", {})
|
|
|
|
# Finance tool should no longer be visible (back to global disable)
|
|
tools_after_clear = await client.list_tools()
|
|
assert not any(t.name == "finance_tool" for t in tools_after_clear)
|
|
|
|
async def test_multiple_rules_accumulate(self):
|
|
"""Test that multiple enable/disable calls accumulate rules."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool(tags={"admin"})
|
|
def admin_tool() -> str:
|
|
return "admin"
|
|
|
|
@mcp.tool
|
|
async def activate_multiple(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
await ctx.enable_components(tags={"admin"})
|
|
rules = await ctx._get_visibility_rules()
|
|
assert len(rules) == 2
|
|
return "activated"
|
|
|
|
# Globally disable finance and admin tools
|
|
mcp.disable(tags={"finance", "admin"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Activate both
|
|
await client.call_tool("activate_multiple", {})
|
|
|
|
# Both should be visible
|
|
tools = await client.list_tools()
|
|
assert any(t.name == "finance_tool" for t in tools)
|
|
assert any(t.name == "admin_tool" for t in tools)
|
|
|
|
async def test_later_rules_override_earlier_rules(self):
|
|
"""Test that later session rules override earlier ones (mark semantics)."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"test"})
|
|
def test_tool() -> str:
|
|
return "test"
|
|
|
|
@mcp.tool
|
|
async def toggle_test(ctx: Context) -> str:
|
|
# First enable, then disable
|
|
await ctx.enable_components(tags={"test"})
|
|
await ctx.disable_components(tags={"test"})
|
|
return "toggled"
|
|
|
|
async with Client(mcp) as client:
|
|
# Toggle (enable then disable)
|
|
await client.call_tool("toggle_test", {})
|
|
|
|
# The disable should win (later mark overrides earlier)
|
|
tools = await client.list_tools()
|
|
assert not any(t.name == "test_tool" for t in tools)
|
|
|
|
async def test_session_transforms_apply_to_resources(self):
|
|
"""Test that session transforms apply to resources too."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.resource("resource://finance", tags={"finance"})
|
|
def finance_resource() -> str:
|
|
return "finance data"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
# Globally disable finance resources
|
|
mcp.disable(tags={"finance"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Before activation, finance resource should not be visible
|
|
resources_before = await client.list_resources()
|
|
assert not any(str(r.uri) == "resource://finance" for r in resources_before)
|
|
|
|
# Activate finance for this session
|
|
await client.call_tool("activate_finance", {})
|
|
|
|
# After activation, finance resource should be visible
|
|
resources_after = await client.list_resources()
|
|
assert any(str(r.uri) == "resource://finance" for r in resources_after)
|
|
|
|
async def test_session_transforms_apply_to_prompts(self):
|
|
"""Test that session transforms apply to prompts too."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.prompt(tags={"finance"})
|
|
def finance_prompt() -> str:
|
|
return "finance prompt"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
# Globally disable finance prompts
|
|
mcp.disable(tags={"finance"})
|
|
|
|
async with Client(mcp) as client:
|
|
# Before activation, finance prompt should not be visible
|
|
prompts_before = await client.list_prompts()
|
|
assert not any(p.name == "finance_prompt" for p in prompts_before)
|
|
|
|
# Activate finance for this session
|
|
await client.call_tool("activate_finance", {})
|
|
|
|
# After activation, finance prompt should be visible
|
|
prompts_after = await client.list_prompts()
|
|
assert any(p.name == "finance_prompt" for p in prompts_after)
|
|
|
|
|
|
class TestSessionVisibilityNotifications:
|
|
"""Test that notifications are sent when session visibility changes."""
|
|
|
|
async def test_enable_components_sends_notifications(self):
|
|
"""Test that enable_components sends all three notification types."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def activate(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
handler = RecordingMessageHandler()
|
|
async with Client(mcp, message_handler=handler) as client:
|
|
handler.reset()
|
|
await client.call_tool("activate", {})
|
|
|
|
# Should receive all three notifications
|
|
tool_notifications = handler.get_notifications(
|
|
"notifications/tools/list_changed"
|
|
)
|
|
resource_notifications = handler.get_notifications(
|
|
"notifications/resources/list_changed"
|
|
)
|
|
prompt_notifications = handler.get_notifications(
|
|
"notifications/prompts/list_changed"
|
|
)
|
|
assert len(tool_notifications) == 1
|
|
assert len(resource_notifications) == 1
|
|
assert len(prompt_notifications) == 1
|
|
|
|
async def test_disable_components_sends_notifications(self):
|
|
"""Test that disable_components sends all three notification types."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def deactivate(ctx: Context) -> str:
|
|
await ctx.disable_components(tags={"finance"})
|
|
return "deactivated"
|
|
|
|
handler = RecordingMessageHandler()
|
|
async with Client(mcp, message_handler=handler) as client:
|
|
handler.reset()
|
|
await client.call_tool("deactivate", {})
|
|
|
|
# Should receive all three notifications
|
|
assert (
|
|
len(handler.get_notifications("notifications/tools/list_changed")) == 1
|
|
)
|
|
assert (
|
|
len(handler.get_notifications("notifications/resources/list_changed"))
|
|
== 1
|
|
)
|
|
assert (
|
|
len(handler.get_notifications("notifications/prompts/list_changed"))
|
|
== 1
|
|
)
|
|
|
|
async def test_clear_visibility_rules_sends_notifications(self):
|
|
"""Test that reset_visibility sends notifications."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def clear(ctx: Context) -> str:
|
|
await ctx.reset_visibility()
|
|
return "cleared"
|
|
|
|
handler = RecordingMessageHandler()
|
|
async with Client(mcp, message_handler=handler) as client:
|
|
handler.reset()
|
|
await client.call_tool("clear", {})
|
|
|
|
# Should receive all three notifications
|
|
assert (
|
|
len(handler.get_notifications("notifications/tools/list_changed")) == 1
|
|
)
|
|
assert (
|
|
len(handler.get_notifications("notifications/resources/list_changed"))
|
|
== 1
|
|
)
|
|
assert (
|
|
len(handler.get_notifications("notifications/prompts/list_changed"))
|
|
== 1
|
|
)
|
|
|
|
async def test_components_hint_limits_notifications(self):
|
|
"""Test that the components hint limits which notifications are sent."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def activate_tools_only(ctx: Context) -> str:
|
|
# Only specify tool components - should only send tool notification
|
|
await ctx.enable_components(tags={"finance"}, components={"tool"})
|
|
return "activated"
|
|
|
|
handler = RecordingMessageHandler()
|
|
async with Client(mcp, message_handler=handler) as client:
|
|
handler.reset()
|
|
await client.call_tool("activate_tools_only", {})
|
|
|
|
# Should only receive tool notification
|
|
assert (
|
|
len(handler.get_notifications("notifications/tools/list_changed")) == 1
|
|
)
|
|
assert (
|
|
len(handler.get_notifications("notifications/resources/list_changed"))
|
|
== 0
|
|
)
|
|
assert (
|
|
len(handler.get_notifications("notifications/prompts/list_changed"))
|
|
== 0
|
|
)
|
|
|
|
|
|
class TestConcurrentSessionIsolation:
|
|
"""Test that concurrent sessions don't leak visibility transforms."""
|
|
|
|
async def test_concurrent_sessions_isolated(self):
|
|
"""Test that two concurrent clients don't leak session transforms."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"finance"})
|
|
def finance_tool() -> str:
|
|
return "finance"
|
|
|
|
@mcp.tool
|
|
async def activate_finance(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"finance"})
|
|
return "activated"
|
|
|
|
# Globally disable finance tools
|
|
mcp.disable(tags={"finance"})
|
|
|
|
# Track what each session sees
|
|
session_a_sees_finance = False
|
|
session_b_sees_finance = False
|
|
ready_event = anyio.Event()
|
|
|
|
async def session_a():
|
|
nonlocal session_a_sees_finance
|
|
async with Client(mcp) as client:
|
|
# Activate finance for this session
|
|
await client.call_tool("activate_finance", {})
|
|
|
|
# Signal that session A has activated
|
|
ready_event.set()
|
|
|
|
# Check that session A sees finance tool
|
|
tools = await client.list_tools()
|
|
session_a_sees_finance = any(t.name == "finance_tool" for t in tools)
|
|
|
|
# Keep session A alive while session B checks
|
|
await anyio.sleep(0.2)
|
|
|
|
async def session_b():
|
|
nonlocal session_b_sees_finance
|
|
# Wait for session A to activate
|
|
await ready_event.wait()
|
|
|
|
async with Client(mcp) as client:
|
|
# Session B should NOT see finance tool
|
|
tools = await client.list_tools()
|
|
session_b_sees_finance = any(t.name == "finance_tool" for t in tools)
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(session_a)
|
|
tg.start_soon(session_b)
|
|
|
|
# Session A should see finance, session B should not
|
|
assert session_a_sees_finance is True, "Session A should see finance tool"
|
|
assert session_b_sees_finance is False, "Session B should NOT see finance tool"
|
|
|
|
async def test_many_concurrent_sessions_isolated(self):
|
|
"""Test that many concurrent sessions remain properly isolated."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"premium"})
|
|
def premium_tool() -> str:
|
|
return "premium"
|
|
|
|
@mcp.tool
|
|
async def activate_premium(ctx: Context) -> str:
|
|
await ctx.enable_components(tags={"premium"})
|
|
return "activated"
|
|
|
|
# Globally disable premium tools
|
|
mcp.disable(tags={"premium"})
|
|
|
|
results: dict[str, bool] = {}
|
|
|
|
async def activated_session(session_id: str):
|
|
async with Client(mcp) as client:
|
|
await client.call_tool("activate_premium", {})
|
|
tools = await client.list_tools()
|
|
results[session_id] = any(t.name == "premium_tool" for t in tools)
|
|
|
|
async def non_activated_session(session_id: str):
|
|
async with Client(mcp) as client:
|
|
tools = await client.list_tools()
|
|
results[session_id] = any(t.name == "premium_tool" for t in tools)
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
# Start 5 activated sessions
|
|
for i in range(5):
|
|
tg.start_soon(activated_session, f"activated_{i}")
|
|
# Start 5 non-activated sessions
|
|
for i in range(5):
|
|
tg.start_soon(non_activated_session, f"non_activated_{i}")
|
|
|
|
# All activated sessions should see premium tool
|
|
for i in range(5):
|
|
assert results[f"activated_{i}"] is True, (
|
|
f"Activated session {i} should see premium tool"
|
|
)
|
|
|
|
# All non-activated sessions should NOT see premium tool
|
|
for i in range(5):
|
|
assert results[f"non_activated_{i}"] is False, (
|
|
f"Non-activated session {i} should NOT see premium tool"
|
|
)
|
|
|
|
|
|
class TestSessionVisibilityResetBug:
|
|
"""Regression tests for #3034: visibility marks leak via shared component mutation."""
|
|
|
|
async def test_disable_then_reset_restores_tools(self):
|
|
"""After disable + reset within the same session, tools should reappear."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"system"})
|
|
def my_tool() -> str:
|
|
return "hello"
|
|
|
|
@mcp.tool(tags={"env"})
|
|
async def enter_env(ctx: Context) -> str:
|
|
await ctx.disable_components(tags={"system"})
|
|
return "entered"
|
|
|
|
@mcp.tool(tags={"env"})
|
|
async def exit_env(ctx: Context) -> str:
|
|
await ctx.reset_visibility()
|
|
return "exited"
|
|
|
|
async with Client(mcp) as client:
|
|
# Tool visible initially
|
|
tools = await client.list_tools()
|
|
assert any(t.name == "my_tool" for t in tools)
|
|
|
|
# Disable it
|
|
await client.call_tool("enter_env", {})
|
|
tools = await client.list_tools()
|
|
assert not any(t.name == "my_tool" for t in tools)
|
|
|
|
# Reset — tool should come back
|
|
await client.call_tool("exit_env", {})
|
|
tools = await client.list_tools()
|
|
assert any(t.name == "my_tool" for t in tools), (
|
|
"Tool should be visible again after reset_visibility"
|
|
)
|
|
|
|
async def test_disable_reset_loop(self):
|
|
"""Repeated disable/reset cycles should work every time (the exact bug from #3034)."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"system"})
|
|
def create_project() -> str:
|
|
return "created"
|
|
|
|
@mcp.tool(tags={"env"})
|
|
async def enter_env(ctx: Context) -> str:
|
|
await ctx.disable_components(tags={"system"})
|
|
return "entered"
|
|
|
|
@mcp.tool(tags={"env"})
|
|
async def exit_env(ctx: Context) -> str:
|
|
await ctx.reset_visibility()
|
|
return "exited"
|
|
|
|
async with Client(mcp) as client:
|
|
for i in range(3):
|
|
# create_project should be visible
|
|
tools = await client.list_tools()
|
|
assert any(t.name == "create_project" for t in tools), (
|
|
f"Iteration {i}: create_project should be visible before enter_env"
|
|
)
|
|
|
|
# Enter env — disables system tools
|
|
await client.call_tool("enter_env", {})
|
|
tools = await client.list_tools()
|
|
assert not any(t.name == "create_project" for t in tools), (
|
|
f"Iteration {i}: create_project should be hidden after enter_env"
|
|
)
|
|
|
|
# Exit env — reset
|
|
await client.call_tool("exit_env", {})
|
|
|
|
async def test_session_disable_does_not_leak_to_concurrent_session(self):
|
|
"""Disabling tools in one session must not affect a concurrent session."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"system"})
|
|
def shared_tool() -> str:
|
|
return "shared"
|
|
|
|
@mcp.tool
|
|
async def disable_system(ctx: Context) -> str:
|
|
await ctx.disable_components(tags={"system"})
|
|
return "disabled"
|
|
|
|
session_b_sees_tool = False
|
|
ready = anyio.Event()
|
|
check_done = anyio.Event()
|
|
|
|
async def session_a():
|
|
async with Client(mcp) as client:
|
|
await client.call_tool("disable_system", {})
|
|
ready.set()
|
|
await check_done.wait()
|
|
|
|
async def session_b():
|
|
nonlocal session_b_sees_tool
|
|
await ready.wait()
|
|
async with Client(mcp) as client:
|
|
tools = await client.list_tools()
|
|
session_b_sees_tool = any(t.name == "shared_tool" for t in tools)
|
|
check_done.set()
|
|
|
|
async with anyio.create_task_group() as tg:
|
|
tg.start_soon(session_a)
|
|
tg.start_soon(session_b)
|
|
|
|
assert session_b_sees_tool is True, (
|
|
"Session B should still see shared_tool despite Session A disabling it"
|
|
)
|
|
|
|
async def test_session_disable_does_not_leak_to_sequential_session(self):
|
|
"""Disabling tools in one session must not affect a later session."""
|
|
from fastmcp import Client
|
|
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool(tags={"system"})
|
|
def shared_tool() -> str:
|
|
return "shared"
|
|
|
|
@mcp.tool
|
|
async def disable_system(ctx: Context) -> str:
|
|
await ctx.disable_components(tags={"system"})
|
|
return "disabled"
|
|
|
|
# Session A disables the tool (no reset)
|
|
async with Client(mcp) as client_a:
|
|
await client_a.call_tool("disable_system", {})
|
|
tools = await client_a.list_tools()
|
|
assert not any(t.name == "shared_tool" for t in tools)
|
|
|
|
# Session B should see it fresh
|
|
async with Client(mcp) as client_b:
|
|
tools = await client_b.list_tools()
|
|
assert any(t.name == "shared_tool" for t in tools), (
|
|
"New session should see shared_tool regardless of previous session"
|
|
)
|