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