prefecthq--fastmcp/tests/server/test_context.py

from unittest.mock import MagicMock

import pytest
from mcp_types import ModelPreferences

from fastmcp.server.context import (
    Context,
    reset_transport,
    set_transport,
)
from fastmcp.server.sampling.run import _parse_model_preferences
from fastmcp.server.server import FastMCP


@pytest.fixture
def context():
    return Context(fastmcp=FastMCP())


class TestParseModelPreferences:
    def test_parse_model_preferences_string(self, context):
        mp = _parse_model_preferences("claude-haiku-4-5")
        assert isinstance(mp, ModelPreferences)
        assert mp.hints is not None
        assert mp.hints[0].name == "claude-haiku-4-5"

    def test_parse_model_preferences_list(self, context):
        mp = _parse_model_preferences(["claude-haiku-4-5", "claude"])
        assert isinstance(mp, ModelPreferences)
        assert mp.hints is not None
        assert [h.name for h in mp.hints] == ["claude-haiku-4-5", "claude"]

    def test_parse_model_preferences_object(self, context):
        obj = ModelPreferences(hints=[])
        assert _parse_model_preferences(obj) is obj

    def test_parse_model_preferences_invalid_type(self, context):
        with pytest.raises(ValueError):
            _parse_model_preferences(model_preferences=123)  # pyright: ignore[reportArgumentType] # type: ignore[invalid-argument-type]  # ty:ignore[invalid-argument-type]


class TestSessionId:
    def test_session_id_with_http_headers(self, context):
        """Test that session_id returns the value from mcp-session-id header."""
        from fastmcp.server.dependencies import (
            FastMCPRequestContext,
            fastmcp_request_ctx,
        )

        mock_headers = {"mcp-session-id": "test-session-123"}

        token = fastmcp_request_ctx.set(
            FastMCPRequestContext(
                session=MagicMock(wraps={}),
                request_id="0",
                meta=None,
                request=MagicMock(headers=mock_headers),
                protocol_version="2025-06-18",
                close_sse_stream=None,
                lifespan_context=MagicMock(),
                _srctx=MagicMock(meta=None),
            )
        )

        try:
            assert context.session_id == "test-session-123"
        finally:
            fastmcp_request_ctx.reset(token)

    def test_session_id_without_http_headers(self, context):
        """Test that session_id returns a UUID when no HTTP headers are available.

        For STDIO/SSE/in-memory transports, we generate a UUID and cache it
        on the session for consistency with state operations.
        """
        import uuid

        from fastmcp.server.dependencies import (
            FastMCPRequestContext,
            fastmcp_request_ctx,
        )

        mock_session = MagicMock(wraps={})
        token = fastmcp_request_ctx.set(
            FastMCPRequestContext(
                session=mock_session,
                request_id="0",
                meta=None,
                request=None,
                protocol_version="2025-06-18",
                close_sse_stream=None,
                lifespan_context=MagicMock(),
                _srctx=MagicMock(meta=None),
            )
        )

        try:
            # session_id should be a valid UUID for non-HTTP transports
            session_id = context.session_id
            assert uuid.UUID(session_id)  # Valid UUID format
            # Should be cached on session
            assert mock_session._fastmcp_state_prefix == session_id
        finally:
            fastmcp_request_ctx.reset(token)


class TestContextState:
    """Test suite for Context state functionality."""

    async def test_context_state_basic(self):
        """Test basic get/set/delete state operations."""
        server = FastMCP("test")
        mock_session = MagicMock()  # Use same session for consistent id()

        async with Context(fastmcp=server, session=mock_session) as context:
            # Initially empty
            assert await context.get_state("test1") is None
            assert await context.get_state("test2") is None

            # Set values
            await context.set_state("test1", "value")
            await context.set_state("test2", 2)

            # Retrieve values
            assert await context.get_state("test1") == "value"
            assert await context.get_state("test2") == 2

            # Update value
            await context.set_state("test1", "new_value")
            assert await context.get_state("test1") == "new_value"

            # Delete value
            await context.delete_state("test1")
            assert await context.get_state("test1") is None

    async def test_context_state_session_isolation(self):
        """Test that different sessions have isolated state."""
        server = FastMCP("test")
        session_a = MagicMock()
        session_b = MagicMock()

        async with Context(fastmcp=server, session=session_a) as context1:
            await context1.set_state("key", "value-from-A")

        async with Context(fastmcp=server, session=session_b) as context2:
            # Session B should not see session A's state
            assert await context2.get_state("key") is None
            await context2.set_state("key", "value-from-B")
            assert await context2.get_state("key") == "value-from-B"

        # Verify session A's state is still intact
        async with Context(fastmcp=server, session=session_a) as context3:
            assert await context3.get_state("key") == "value-from-A"

    async def test_context_state_persists_across_requests(self):
        """Test that state persists across multiple context instances (requests)."""
        server = FastMCP("test")
        mock_session = MagicMock()  # Same session = same id()

        # First request sets state
        async with Context(fastmcp=server, session=mock_session) as context1:
            await context1.set_state("counter", 1)

        # Second request in same session sees the state
        async with Context(fastmcp=server, session=mock_session) as context2:
            counter = await context2.get_state("counter")
            assert counter == 1
            await context2.set_state("counter", counter + 1)

        # Third request sees updated state
        async with Context(fastmcp=server, session=mock_session) as context3:
            assert await context3.get_state("counter") == 2

    async def test_context_state_nested_contexts_share_state(self):
        """Test that nested contexts within the same session share state."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context1:
            await context1.set_state("key", "outer-value")

            async with Context(fastmcp=server, session=mock_session) as context2:
                # Nested context sees same state (same session)
                assert await context2.get_state("key") == "outer-value"

                # Nested context can modify shared state
                await context2.set_state("key", "inner-value")

            # Outer context sees the modification
            assert await context1.get_state("key") == "inner-value"

    async def test_two_clients_same_key_isolated_by_session(self):
        """Test that two different clients can store the same key independently.

        Each client gets an auto-generated session ID, and their state is isolated.
        """
        import json

        from fastmcp import Client

        server = FastMCP("test")
        stored_session_ids: list[str] = []

        @server.tool
        async def store_and_read(value: str, ctx: Context) -> dict:
            """Store a value and return all state info."""
            stored_session_ids.append(ctx.session_id)
            existing = await ctx.get_state("shared_key")
            await ctx.set_state("shared_key", value)
            new_value = await ctx.get_state("shared_key")
            return {
                "session_id": ctx.session_id,
                "existing_value": existing,
                "new_value": new_value,
            }

        # Client 1 stores "value-from-client-1"
        async with Client(server) as client1:
            result1 = await client1.call_tool(
                "store_and_read", {"value": "value-from-client-1"}
            )
            data1 = json.loads(result1.content[0].text)
            assert data1["existing_value"] is None  # First write
            assert data1["new_value"] == "value-from-client-1"
            session_id_1 = data1["session_id"]

        # Client 2 stores "value-from-client-2" with the SAME key
        async with Client(server) as client2:
            result2 = await client2.call_tool(
                "store_and_read", {"value": "value-from-client-2"}
            )
            data2 = json.loads(result2.content[0].text)
            # Client 2 should NOT see client 1's value (different session)
            assert data2["existing_value"] is None
            assert data2["new_value"] == "value-from-client-2"
            session_id_2 = data2["session_id"]

        # Verify session IDs were auto-generated and are different
        assert session_id_1 is not None
        assert session_id_2 is not None
        assert session_id_1 != session_id_2

        # Client 1 reconnects and should still see their value
        async with Client(server) as client1_again:
            # But this is a NEW session (new connection = new session ID)
            result3 = await client1_again.call_tool(
                "store_and_read", {"value": "value-from-client-1-again"}
            )
            data3 = json.loads(result3.content[0].text)
            # New session, so existing value is None
            assert data3["existing_value"] is None
            assert data3["session_id"] != session_id_1  # Different session


class TestContextStateSerializable:
    """Tests for the serializable parameter on set_state."""

    async def test_set_state_serializable_false_stores_arbitrary_objects(self):
        """Non-serializable objects can be stored with serializable=False."""
        server = FastMCP("test")
        mock_session = MagicMock()

        class MyClient:
            def __init__(self):
                self.connected = True

        client = MyClient()

        async with Context(fastmcp=server, session=mock_session) as context:
            await context.set_state("client", client, serializable=False)
            result = await context.get_state("client")
            assert result is client
            assert result.connected is True

    async def test_set_state_serializable_false_does_not_persist_across_requests(self):
        """Non-serializable state is request-scoped and gone in a new context."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context:
            await context.set_state("key", object(), serializable=False)
            assert await context.get_state("key") is not None

        async with Context(fastmcp=server, session=mock_session) as context:
            assert await context.get_state("key") is None

    async def test_set_state_serializable_true_rejects_non_serializable(self):
        """Default set_state raises TypeError for non-serializable values."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context:
            with pytest.raises(TypeError, match="serializable=False"):
                await context.set_state("key", object())

    async def test_set_state_serializable_false_shadows_session_state(self):
        """Request-scoped state shadows session-scoped state for the same key."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context:
            await context.set_state("key", "session-value")
            assert await context.get_state("key") == "session-value"

            await context.set_state("key", "request-value", serializable=False)
            assert await context.get_state("key") == "request-value"

    async def test_delete_state_removes_from_both_stores(self):
        """delete_state clears both request-scoped and session-scoped values."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context:
            await context.set_state("key", "session-value")
            await context.set_state("key", "request-value", serializable=False)
            assert await context.get_state("key") == "request-value"

            await context.delete_state("key")
            assert await context.get_state("key") is None

    async def test_serializable_state_still_persists_across_requests(self):
        """Serializable state (default) still persists across requests."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context:
            await context.set_state("key", "persistent")

        async with Context(fastmcp=server, session=mock_session) as context:
            assert await context.get_state("key") == "persistent"

    async def test_serializable_write_clears_request_scoped_shadow(self):
        """Writing serializable state clears any request-scoped shadow for the same key."""
        server = FastMCP("test")
        mock_session = MagicMock()

        async with Context(fastmcp=server, session=mock_session) as context:
            await context.set_state("key", "request-value", serializable=False)
            assert await context.get_state("key") == "request-value"

            # Serializable write should clear the shadow
            await context.set_state("key", "session-value")
            assert await context.get_state("key") == "session-value"


class TestContextMeta:
    """Test suite for Context meta functionality."""

    def test_request_context_meta_access(self, context):
        """Test that the lifted _meta dict is accessible from request context.

        In the v2 SDK port the wrapper's ``meta`` is the raw ``_meta`` block
        lifted from the request params (a dict), not a typed object.
        """
        from fastmcp.server.dependencies import (
            FastMCPRequestContext,
            fastmcp_request_ctx,
        )

        meta_dict = {
            "user_id": "user-123",
            "trace_id": "trace-456",
            "custom_field": "custom-value",
        }

        token = fastmcp_request_ctx.set(
            FastMCPRequestContext(
                session=MagicMock(wraps={}),
                request_id="0",
                meta=meta_dict,
                request=None,
                protocol_version="2025-06-18",
                close_sse_stream=None,
                lifespan_context=MagicMock(),
                _srctx=MagicMock(meta=None),
            )
        )

        retrieved_meta = context.request_context.meta
        assert retrieved_meta is not None
        assert retrieved_meta["user_id"] == "user-123"
        assert retrieved_meta["trace_id"] == "trace-456"
        assert retrieved_meta["custom_field"] == "custom-value"

        fastmcp_request_ctx.reset(token)

    def test_request_context_meta_none(self, context):
        """Test that context handles None meta gracefully."""
        from fastmcp.server.dependencies import (
            FastMCPRequestContext,
            fastmcp_request_ctx,
        )

        token = fastmcp_request_ctx.set(
            FastMCPRequestContext(
                session=MagicMock(wraps={}),
                request_id="0",
                meta=None,
                request=None,
                protocol_version="2025-06-18",
                close_sse_stream=None,
                lifespan_context=MagicMock(),
                _srctx=MagicMock(meta=None),
            )
        )

        retrieved_meta = context.request_context.meta
        assert retrieved_meta is None

        fastmcp_request_ctx.reset(token)


class TestTransport:
    """Test suite for Context transport property."""

    def test_transport_returns_none_outside_server_context(self, context):
        """Test that transport returns None when not in a server context."""
        assert context.transport is None

    def test_transport_returns_stdio(self, context):
        """Test that transport returns 'stdio' when set."""
        token = set_transport("stdio")
        try:
            assert context.transport == "stdio"
        finally:
            reset_transport(token)

    def test_transport_returns_sse(self, context):
        """Test that transport returns 'sse' when set."""
        token = set_transport("sse")
        try:
            assert context.transport == "sse"
        finally:
            reset_transport(token)

    def test_transport_returns_streamable_http(self, context):
        """Test that transport returns 'streamable-http' when set."""
        token = set_transport("streamable-http")
        try:
            assert context.transport == "streamable-http"
        finally:
            reset_transport(token)

    def test_transport_reset(self, context):
        """Test that transport resets correctly."""
        assert context.transport is None
        token = set_transport("stdio")
        assert context.transport == "stdio"
        reset_transport(token)
        assert context.transport is None


class TestTransportIntegration:
    """Integration tests for transport property with actual server/client."""

    async def test_transport_in_tool_via_client(self):
        """Test that transport is accessible from within a tool via Client."""
        from fastmcp import Client

        mcp = FastMCP("test")
        observed_transport = None

        @mcp.tool
        def get_transport(ctx: Context) -> str:
            nonlocal observed_transport
            observed_transport = ctx.transport
            return observed_transport or "none"

        # Client uses in-memory transport which doesn't set transport type
        # so we expect None here (the transport is only set by run_* methods)
        async with Client(mcp) as client:
            result = await client.call_tool("get_transport", {})
            assert observed_transport is None
            assert result.data == "none"

    async def test_transport_set_manually_is_visible_in_tool(self):
        """Test that manually set transport is visible from within a tool."""
        from fastmcp import Client

        mcp = FastMCP("test")
        observed_transport = None

        @mcp.tool
        def get_transport(ctx: Context) -> str:
            nonlocal observed_transport
            observed_transport = ctx.transport
            return observed_transport or "none"

        # Manually set transport before running
        token = set_transport("stdio")
        try:
            async with Client(mcp) as client:
                result = await client.call_tool("get_transport", {})
                assert observed_transport == "stdio"
                assert result.data == "stdio"
        finally:
            reset_transport(token)

    async def test_transport_set_via_http_middleware(self):
        """Test that transport is set per-request via HTTP middleware."""
        from fastmcp import Client
        from fastmcp.client.transports import StreamableHttpTransport
        from fastmcp.utilities.tests import run_server_async

        mcp = FastMCP("test")
        observed_transport = None

        @mcp.tool
        def get_transport(ctx: Context) -> str:
            nonlocal observed_transport
            observed_transport = ctx.transport
            return observed_transport or "none"

        async with run_server_async(mcp, transport="streamable-http") as url:
            transport = StreamableHttpTransport(url=url)
            async with Client(transport=transport) as client:
                result = await client.call_tool("get_transport", {})
                assert observed_transport == "streamable-http"
                assert result.data == "streamable-http"