core/tests/components/mcp_server/test_http.py

"""Test the Model Context Protocol Server init module."""

from collections.abc import AsyncGenerator, Callable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from http import HTTPStatus
import json
import logging
from typing import Any
from unittest.mock import AsyncMock, patch

import aiohttp
import mcp
import mcp.client.session
import mcp.client.sse
import mcp.client.streamable_http
from mcp.shared.exceptions import McpError
import probatio
import pytest

from homeassistant.components.conversation import DOMAIN as CONVERSATION_DOMAIN
from homeassistant.components.homeassistant.exposed_entities import async_expose_entity
from homeassistant.components.intent import async_register_timer_handler
from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
from homeassistant.components.mcp_server.const import DOMAIN, STATELESS_LLM_API
from homeassistant.components.mcp_server.http import (
    MESSAGES_API,
    SSE_API,
    STREAMABLE_API,
)
from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import (
    CONF_LLM_HASS_API,
    CONTENT_TYPE_JSON,
    STATE_OFF,
    STATE_ON,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers import (
    area_registry as ar,
    device_registry as dr,
    entity_registry as er,
    llm,
)
from homeassistant.helpers.httpx_client import create_async_httpx_client
from homeassistant.setup import async_setup_component
from homeassistant.util.json import JsonObjectType

from .conftest import TEST_LLM_API_ID, MockLLMAPI

from tests.common import MockConfigEntry, setup_test_component_platform
from tests.components.light.common import MockLight
from tests.typing import ClientSessionGenerator

_LOGGER = logging.getLogger(__name__)

TEST_ENTITY = "light.kitchen"
DEVICE_ID_META_KEY = "io.home-assistant/device_id"
SNAPSHOT_RESOURCE_URI = "homeassistant://assist/context-snapshot"
type MCPClientFactory = Callable[
    [HomeAssistant, str, str],
    AbstractAsyncContextManager[mcp.client.session.ClientSession],
]
INITIALIZE_MESSAGE = {
    "jsonrpc": "2.0",
    "id": "request-id-1",
    "method": "initialize",
    "params": {
        "protocolVersion": "1.0",
        "capabilities": {},
        "clientInfo": {
            "name": "test",
            "version": "1",
        },
    },
}
EVENT_PREFIX = "event: "
DATA_PREFIX = "data: "
EXPECTED_PROMPT_ENTITY_DEFINITION = """
- names: Kitchen Light
  domain: light
  areas: Kitchen
"""


class _StubTool(llm.Tool):
    """Minimal tool with a configurable parameter schema."""

    name = "test_tool"

    def __init__(self, parameters: probatio.Schema) -> None:
        """Initialize the stub tool."""
        self.parameters = parameters

    async def async_call(
        self,
        hass: HomeAssistant,
        tool_input: llm.ToolInput,
        llm_context: llm.LLMContext,
    ) -> JsonObjectType:
        """Return an empty result."""
        return {}


@pytest.fixture
async def setup_integration(hass: HomeAssistant, config_entry: MockConfigEntry) -> None:
    """Set up the config entry."""
    await hass.config_entries.async_setup(config_entry.entry_id)
    assert config_entry.state is ConfigEntryState.LOADED


@pytest.fixture(autouse=True)
async def mock_entities(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
    area_registry: ar.AreaRegistry,
    setup_integration: None,
) -> None:
    """Fixture to expose entities to the conversation agent."""
    entity = MockLight("Kitchen Light", STATE_OFF)
    entity.entity_id = TEST_ENTITY
    entity.unique_id = "test-light-unique-id"
    setup_test_component_platform(hass, LIGHT_DOMAIN, [entity])

    assert await async_setup_component(
        hass,
        LIGHT_DOMAIN,
        {LIGHT_DOMAIN: [{"platform": "test"}]},
    )
    await hass.async_block_till_done()
    kitchen = area_registry.async_get_or_create("Kitchen")
    entity_registry.async_update_entity(TEST_ENTITY, area_id=kitchen.id)

    async_expose_entity(hass, CONVERSATION_DOMAIN, TEST_ENTITY, True)


async def sse_response_reader(
    response: aiohttp.ClientResponse,
) -> AsyncGenerator[tuple[str, str]]:
    """Read SSE responses from the server and emit event messages.

    SSE responses are formatted as:
        event: event-name
        data: event-data
    and this function emits each event-name and event-data as a tuple.
    """
    it = aiter(response.content)
    while True:
        line = (await anext(it)).decode()
        if not line.startswith(EVENT_PREFIX):
            raise ValueError("Expected event")
        event = line[len(EVENT_PREFIX) :].strip()
        line = (await anext(it)).decode()
        if not line.startswith(DATA_PREFIX):
            raise ValueError("Expected data")
        data = line[len(DATA_PREFIX) :].strip()
        line = (await anext(it)).decode()
        assert line == "\r\n"
        yield event, data


async def test_http_sse(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test SSE endpoint can be used to receive MCP messages."""

    client = await hass_client()

    # Start an SSE session
    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.OK

    # Decode a single SSE response that sends the messages endpoint
    reader = sse_response_reader(response)
    event, endpoint_url = await anext(reader)
    assert event == "endpoint"

    # Send an initialize message on the messages endpoint
    response = await client.post(endpoint_url, json=INITIALIZE_MESSAGE)
    assert response.status == HTTPStatus.OK

    # Decode the initialize response event message from the SSE stream
    event, data = await anext(reader)
    assert event == "message"
    message = json.loads(data)
    assert message.get("jsonrpc") == "2.0"
    assert message.get("id") == "request-id-1"
    assert "serverInfo" in message.get("result", {})
    assert "protocolVersion" in message.get("result", {})


async def test_http_messages_missing_session_id(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test the tools list endpoint."""

    client = await hass_client()
    response = await client.post(MESSAGES_API.format(session_id="invalid-session-id"))
    assert response.status == HTTPStatus.NOT_FOUND
    response_data = await response.text()
    assert response_data == "Could not find session ID 'invalid-session-id'"


async def test_http_messages_invalid_message_format(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test the tools list endpoint."""

    client = await hass_client()
    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.OK
    reader = sse_response_reader(response)
    event, endpoint_url = await anext(reader)
    assert event == "endpoint"

    response = await client.post(endpoint_url, json={"invalid": "message"})
    assert response.status == HTTPStatus.BAD_REQUEST
    response_data = await response.text()
    assert response_data == "Could not parse message"


async def test_http_sse_multiple_config_entries(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test the SSE endpoint will fail with multiple config entries.

    This cannot happen in practice as the integration only supports a single
    config entry, but this is added for test coverage.
    """

    config_entry = MockConfigEntry(
        domain=DOMAIN, data={CONF_LLM_HASS_API: ["llm-api-id"]}
    )
    config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(config_entry.entry_id)

    client = await hass_client()

    # Attempt to start an SSE session will fail
    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.NOT_FOUND
    response_data = await response.text()
    assert "Found multiple Model Context Protocol" in response_data


async def test_http_sse_no_config_entry(
    hass: HomeAssistant,
    setup_integration: None,
    config_entry: MockConfigEntry,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test the SSE endpoint fails with a missing config entry."""

    await hass.config_entries.async_unload(config_entry.entry_id)
    assert config_entry.state is ConfigEntryState.NOT_LOADED

    client = await hass_client()

    # Start an SSE session
    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.NOT_FOUND
    response_data = await response.text()
    assert "Model Context Protocol server is not configured" in response_data


async def test_http_messages_no_config_entry(
    hass: HomeAssistant,
    setup_integration: None,
    config_entry: MockConfigEntry,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test the message endpoint will fail if the config entry is unloaded."""

    client = await hass_client()

    # Start an SSE session
    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.OK
    reader = sse_response_reader(response)
    event, endpoint_url = await anext(reader)
    assert event == "endpoint"

    # Invalidate the session by unloading the config entry
    await hass.config_entries.async_unload(config_entry.entry_id)
    assert config_entry.state is ConfigEntryState.NOT_LOADED

    # Reload the config entry and ensure the session is not found
    await hass.config_entries.async_setup(config_entry.entry_id)
    assert config_entry.state is ConfigEntryState.LOADED

    response = await client.post(endpoint_url, json=INITIALIZE_MESSAGE)
    assert response.status == HTTPStatus.NOT_FOUND
    response_data = await response.text()
    assert "Could not find session ID" in response_data


async def test_options_flow_closes_sessions(
    hass: HomeAssistant,
    setup_integration: None,
    config_entry: MockConfigEntry,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test an open SSE session is closed when the selected APIs change."""
    llm.async_register_api(
        hass, MockLLMAPI(hass=hass, id=TEST_LLM_API_ID, name="Test API")
    )
    client = await hass_client()

    # Start an SSE session
    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.OK
    reader = sse_response_reader(response)
    event, endpoint_url = await anext(reader)
    assert event == "endpoint"

    # Change the exposed LLM APIs
    result = await hass.config_entries.options.async_init(config_entry.entry_id)
    await hass.config_entries.options.async_configure(
        result["flow_id"],
        {CONF_LLM_HASS_API: [TEST_LLM_API_ID]},
    )
    await hass.async_block_till_done()

    # The session serving the previous APIs is gone
    response = await client.post(endpoint_url, json=INITIALIZE_MESSAGE)
    assert response.status == HTTPStatus.NOT_FOUND
    response_data = await response.text()
    assert "Could not find session ID" in response_data


async def test_http_requires_authentication(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client_no_auth: ClientSessionGenerator,
) -> None:
    """Test the SSE endpoint requires authentication."""

    client = await hass_client_no_auth()

    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.UNAUTHORIZED

    response = await client.post(MESSAGES_API.format(session_id="session-id"))
    assert response.status == HTTPStatus.UNAUTHORIZED


@pytest.fixture(params=["sse", "streamable"])
def mcp_protocol(request: pytest.FixtureRequest):
    """Fixture to parametrize tests with different MCP protocols."""
    return request.param


@pytest.fixture
async def mcp_url(mcp_protocol: str, hass_client: ClientSessionGenerator) -> str:
    """Fixture to get the MCP integration URL."""
    if mcp_protocol == "sse":
        url = SSE_API
    else:
        url = STREAMABLE_API
    client = await hass_client()
    return str(client.make_url(url))


@asynccontextmanager
async def mcp_sse_session(
    hass: HomeAssistant,
    mcp_url: str,
    hass_supervisor_access_token: str,
) -> AsyncGenerator[mcp.client.session.ClientSession]:
    """Create an MCP session."""

    headers = {"Authorization": f"Bearer {hass_supervisor_access_token}"}

    async with (
        mcp.client.sse.sse_client(mcp_url, headers=headers) as streams,
        mcp.client.session.ClientSession(*streams) as session,
    ):
        await session.initialize()
        yield session


@asynccontextmanager
async def mcp_streamable_session(
    hass: HomeAssistant,
    mcp_url: str,
    hass_supervisor_access_token: str,
) -> AsyncGenerator[mcp.client.session.ClientSession]:
    """Create an MCP session."""

    headers = {"Authorization": f"Bearer {hass_supervisor_access_token}"}

    async with (
        mcp.client.streamable_http.streamable_http_client(
            mcp_url, http_client=create_async_httpx_client(hass, headers=headers)
        ) as (read_stream, write_stream, _),
        mcp.client.session.ClientSession(read_stream, write_stream) as session,
    ):
        await session.initialize()
        yield session


@pytest.fixture(name="mcp_client")
def mcp_client_fixture(mcp_protocol: str) -> Any:
    """Fixture to parametrize tests with different MCP clients."""
    if mcp_protocol == "sse":
        return mcp_sse_session
    if mcp_protocol == "streamable":
        return mcp_streamable_session
    raise ValueError(f"Unknown MCP protocol: {mcp_protocol}")


@pytest.mark.parametrize(
    ("mcp_request", "result_type"),
    [
        pytest.param(
            mcp.types.ClientRequest(mcp.types.ListToolsRequest()),
            mcp.types.ListToolsResult,
            id="tools-list",
        ),
        pytest.param(
            mcp.types.ClientRequest(
                mcp.types.CallToolRequest(
                    params=mcp.types.CallToolRequestParams(
                        name="intent__HassTurnOn",
                        arguments={"name": "kitchen light"},
                    )
                )
            ),
            mcp.types.CallToolResult,
            id="tools-call",
        ),
        pytest.param(
            mcp.types.ClientRequest(mcp.types.ListPromptsRequest()),
            mcp.types.ListPromptsResult,
            id="prompts-list",
        ),
        pytest.param(
            mcp.types.ClientRequest(
                mcp.types.GetPromptRequest(
                    params=mcp.types.GetPromptRequestParams(name="Assist")
                )
            ),
            mcp.types.GetPromptResult,
            id="prompts-get",
        ),
        pytest.param(
            mcp.types.ClientRequest(mcp.types.ListResourcesRequest()),
            mcp.types.ListResourcesResult,
            id="resources-list",
        ),
        pytest.param(
            mcp.types.ClientRequest.model_validate(
                {"method": "resources/read", "params": {"uri": SNAPSHOT_RESOURCE_URI}}
            ),
            mcp.types.ReadResourceResult,
            id="resources-read",
        ),
    ],
)
async def test_request_device_id(
    hass: HomeAssistant,
    mcp_url: str,
    mcp_client: MCPClientFactory,
    hass_supervisor_access_token: str,
    mcp_request: mcp.types.ClientRequest,
    result_type: type[mcp.types.Result],
) -> None:
    """Apply the caller device to each request without retaining it in the session."""
    request_data = mcp_request.model_dump(by_alias=True, exclude_none=True)
    request_data.setdefault("params", {})["_meta"] = {DEVICE_ID_META_KEY: "test-device"}

    with patch(
        "homeassistant.helpers.llm.async_get_api", wraps=llm.async_get_api
    ) as mock_get_api:
        async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
            await session.send_request(
                mcp.types.ClientRequest.model_validate(request_data), result_type
            )
            assert mock_get_api.await_count > 0
            device_contexts = [call.args[2] for call in mock_get_api.await_args_list]
            assert {context.device_id for context in device_contexts} == {"test-device"}
            mock_get_api.reset_mock()

            await session.send_request(mcp_request, result_type)

    assert mock_get_api.await_count > 0
    contexts = [call.args[2] for call in mock_get_api.await_args_list]
    assert {context.device_id for context in contexts} == {None}
    assert {context.device_id for context in device_contexts} == {"test-device"}


@pytest.mark.parametrize(
    "metadata",
    [
        pytest.param({}, id="empty"),
        pytest.param({"other": "value"}, id="unrelated"),
        pytest.param({DEVICE_ID_META_KEY: None}, id="null-device"),
    ],
)
async def test_request_metadata_without_device_id(
    hass: HomeAssistant,
    mcp_url: str,
    mcp_client: MCPClientFactory,
    hass_supervisor_access_token: str,
    metadata: dict[str, str | None],
) -> None:
    """Metadata without a caller device keeps the default LLM context."""
    with patch(
        "homeassistant.helpers.llm.async_get_api", wraps=llm.async_get_api
    ) as mock_get_api:
        async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
            await session.list_tools(
                params=mcp.types.PaginatedRequestParams(
                    _meta=mcp.types.RequestParams.Meta.model_validate(metadata)
                )
            )

    mock_get_api.assert_awaited_once()
    assert mock_get_api.await_args.args[2].device_id is None


@pytest.mark.parametrize(
    "device_id",
    [
        pytest.param(123, id="number"),
        pytest.param(True, id="boolean"),
        pytest.param(["test-device"], id="list"),
        pytest.param({"id": "test-device"}, id="object"),
    ],
)
async def test_request_invalid_device_id(
    hass: HomeAssistant,
    mcp_url: str,
    mcp_client: MCPClientFactory,
    hass_supervisor_access_token: str,
    device_id: int | bool | list[str] | dict[str, str],
) -> None:
    """Reject caller device metadata with an invalid type."""
    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        with pytest.raises(
            McpError, match="io.home-assistant/device_id must be a string"
        ):
            await session.list_tools(
                params=mcp.types.PaginatedRequestParams(
                    _meta=mcp.types.RequestParams.Meta.model_validate(
                        {DEVICE_ID_META_KEY: device_id}
                    )
                )
            )


async def test_tool_call_invalid_device_id(
    hass: HomeAssistant,
    mcp_url: str,
    mcp_client: MCPClientFactory,
    hass_supervisor_access_token: str,
) -> None:
    """Invalid caller metadata returns a tool error without performing the action."""
    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.call_tool(
            name="intent__HassTurnOn",
            arguments={"name": "kitchen light"},
            meta={DEVICE_ID_META_KEY: 123},
        )

    assert result.isError
    assert result.content == [
        mcp.types.TextContent(
            type="text", text="io.home-assistant/device_id must be a string"
        )
    ]
    assert hass.states.get(TEST_ENTITY).state == STATE_OFF


async def test_request_device_id_enables_timer_tools(
    hass: HomeAssistant,
    mcp_url: str,
    mcp_client: MCPClientFactory,
    hass_supervisor_access_token: str,
) -> None:
    """Offer timer tools only for requests from a device that supports timers."""

    def handle_timer(*args: object) -> None:
        pass

    async_register_timer_handler(hass, "test-device", handle_timer)

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.list_tools(
            params=mcp.types.PaginatedRequestParams(
                _meta=mcp.types.RequestParams.Meta.model_validate(
                    {DEVICE_ID_META_KEY: "test-device"}
                )
            )
        )
        assert "intent__HassStartTimer" in {tool.name for tool in result.tools}

        result = await session.list_tools()
        assert "intent__HassStartTimer" not in {tool.name for tool in result.tools}


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_mcp_tools_list(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the tools list endpoint."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.list_tools()

    # Pick a single arbitrary tool and test that description and parameters
    # are converted correctly.
    tool = next(
        iter(tool for tool in result.tools if tool.name == "intent__HassTurnOn")
    )
    assert tool.name == "intent__HassTurnOn"
    assert tool.description is not None
    assert tool.inputSchema
    assert tool.inputSchema.get("type") == "object"
    properties = tool.inputSchema.get("properties")
    assert properties.get("name") == {"type": "string"}


@pytest.mark.parametrize("llm_hass_api", [TEST_LLM_API_ID])
@pytest.mark.parametrize(
    ("parameters", "expected_required"),
    [
        pytest.param(
            probatio.Schema(
                {probatio.Required("name"): str, probatio.Optional("area"): str}
            ),
            ["name"],
            id="required-and-optional",
        ),
        pytest.param(
            probatio.Schema({probatio.Optional("area"): str}),
            None,
            id="optional-only",
        ),
    ],
)
async def test_mcp_tools_list_required_parameters(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: MCPClientFactory,
    hass_supervisor_access_token: str,
    parameters: probatio.Schema,
    expected_required: list[str] | None,
) -> None:
    """Test the tools list advertises the required tool parameters."""

    llm.async_register_api(
        hass,
        MockLLMAPI(
            hass=hass,
            id=TEST_LLM_API_ID,
            name="Test API",
            tools=[_StubTool(parameters)],
        ),
    )

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.list_tools()

    tool = next(iter(tool for tool in result.tools if tool.name == "test_tool"))
    assert tool.inputSchema.get("type") == "object"
    assert tool.inputSchema.get("required") == expected_required


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_mcp_tool_call(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the tool call endpoint."""

    state = hass.states.get("light.kitchen")
    assert state
    assert state.state == STATE_OFF

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.call_tool(
            name="intent__HassTurnOn",
            arguments={"name": "kitchen light"},
        )

    assert not result.isError
    assert len(result.content) == 1
    assert result.content[0].type == "text"
    # The content is the raw tool call payload
    content = json.loads(result.content[0].text)
    assert content.get("data", {}).get("success")
    assert not content.get("data", {}).get("failed")

    # Verify tool call invocation
    state = hass.states.get("light.kitchen")
    assert state
    assert state.state == STATE_ON


async def test_mcp_tool_call_failed(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the tool call endpoint with a failure."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.call_tool(
            name="intent__HassTurnOn",
            arguments={"name": "backyard"},
        )

    assert result.isError
    assert len(result.content) == 1
    assert result.content[0].type == "text"
    assert "Error calling tool" in result.content[0].text


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_prompt_list(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the list prompt endpoint."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.list_prompts()

    assert len(result.prompts) == 1
    prompt = result.prompts[0]
    assert prompt.name == "Assist"
    assert prompt.description == "Default prompt for Home Assistant Assist API"


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_prompt_get(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the get prompt endpoint."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.get_prompt(name="Assist")

    assert result.description == "Default prompt for Home Assistant Assist API"
    assert len(result.messages) == 1
    assert result.messages[0].role == "assistant"
    assert result.messages[0].content.type == "text"
    assert "When controlling Home Assistant" in result.messages[0].content.text
    assert EXPECTED_PROMPT_ENTITY_DEFINITION in result.messages[0].content.text


async def test_get_unknown_prompt(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the get prompt endpoint."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        with pytest.raises(McpError):
            await session.get_prompt(name="Unknown")


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_mcp_resources_list(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the resource list endpoint."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.list_resources()

    assert len(result.resources) == 1
    resource = result.resources[0]
    assert str(resource.uri) == SNAPSHOT_RESOURCE_URI
    assert resource.name == "assist_context_snapshot"
    assert resource.title == "Assist context snapshot"
    assert resource.description is not None
    assert resource.mimeType == "text/plain"


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_mcp_resource_read(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test reading an MCP resource."""

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        resources = await session.list_resources()
        resource = resources.resources[0]
        result = await session.read_resource(resource.uri)

    assert len(result.contents) == 1
    content = result.contents[0]
    assert content.uri == resource.uri
    assert content.mimeType == "text/plain"
    assert content.text == (
        "Live Context: An overview of the areas and the devices in this smart home:\n"
        "- names: Kitchen Light\n"
        "  domain: light\n"
        "  state: 'off'\n"
        "  areas: Kitchen\n"
    )


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST, STATELESS_LLM_API])
async def test_mcp_resource_read_unknown_resource(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test reading an unknown MCP resource."""

    unknown_uri = mcp.types.Resource(
        uri="homeassistant://assist/missing",
        name="missing",
    ).uri

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        with pytest.raises(McpError, match="Unknown resource"):
            await session.read_resource(unknown_uri)


@pytest.mark.parametrize("llm_hass_api", [TEST_LLM_API_ID])
async def test_mcp_resources_unavailable_without_live_context_tool(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test resources are unavailable when the selected API exposes no live context."""

    llm.async_register_api(
        hass, MockLLMAPI(hass=hass, id=TEST_LLM_API_ID, name="Test API")
    )

    resource_uri = mcp.types.Resource(
        uri=SNAPSHOT_RESOURCE_URI,
        name="assist_context_snapshot",
    ).uri

    async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
        result = await session.list_resources()

        assert result.resources == []

        with pytest.raises(McpError, match="Unknown resource"):
            await session.read_resource(resource_uri)


@pytest.mark.parametrize("llm_hass_api", [llm.LLM_API_ASSIST])
async def test_mcp_tool_call_unicode(
    hass: HomeAssistant,
    setup_integration: None,
    mcp_url: str,
    mcp_client: Any,
    hass_supervisor_access_token: str,
) -> None:
    """Test the tool call endpoint preserves unicode characters."""

    # Mock the API instance
    mock_api = AsyncMock()
    mock_api.api.name = "Assist"
    mock_api.tools = []
    mock_api.custom_serializer = None
    mock_api.async_call_tool.return_value = {"message": "这是一个测试"}

    # We need to ensure when the server calls llm.async_get_api, it gets our mock
    # async_get_api is awaited, so we need an AsyncMock
    with patch(
        "homeassistant.helpers.llm.async_get_api", new_callable=AsyncMock
    ) as mock_get_api:
        mock_get_api.return_value = mock_api
        async with mcp_client(hass, mcp_url, hass_supervisor_access_token) as session:
            result = await session.call_tool(
                name="AnyTool",
                arguments={},
            )

    assert not result.isError
    assert len(result.content) == 1
    assert result.content[0].type == "text"

    # Check that the text contains the raw unicode characters, NOT the escaped version
    response_text = result.content[0].text
    assert "这是一个测试" in response_text
    assert "\\u" not in response_text


async def test_streamable_api_id_exposes_registered_api(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
    hass_supervisor_access_token: str,
) -> None:
    """Test the keyed endpoint exposes any registered API, not just the configured one."""
    llm.async_register_api(
        hass, MockLLMAPI(hass=hass, id=TEST_LLM_API_ID, name="Test API")
    )

    client = await hass_client()
    mcp_url = str(client.make_url(f"{STREAMABLE_API}/{TEST_LLM_API_ID}"))

    async with mcp_streamable_session(
        hass, mcp_url, hass_supervisor_access_token
    ) as session:
        result = await session.list_prompts()

    assert len(result.prompts) == 1
    assert result.prompts[0].name == "Test API"


async def test_streamable_api_id_requires_admin(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
    hass_read_only_access_token: str,
) -> None:
    """Test a non-Assist keyed endpoint requires an admin user."""
    llm.async_register_api(
        hass, MockLLMAPI(hass=hass, id=TEST_LLM_API_ID, name="Test API")
    )

    client = await hass_client(hass_read_only_access_token)
    response = await client.post(
        f"{STREAMABLE_API}/{TEST_LLM_API_ID}",
        json=INITIALIZE_MESSAGE,
        headers={"accept": CONTENT_TYPE_JSON},
    )
    assert response.status == HTTPStatus.UNAUTHORIZED


async def test_streamable_api_id_assist_allows_non_admin(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
    hass_read_only_access_token: str,
) -> None:
    """Test the Assist keyed endpoint does not require an admin user."""
    client = await hass_client(hass_read_only_access_token)
    response = await client.post(
        f"{STREAMABLE_API}/{llm.LLM_API_ASSIST}",
        json=INITIALIZE_MESSAGE,
        headers={"accept": CONTENT_TYPE_JSON},
    )
    assert response.status == HTTPStatus.OK


async def test_streamable_api_id_unknown(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test the keyed endpoint returns 404 for an unknown API ID."""
    client = await hass_client()
    response = await client.post(
        f"{STREAMABLE_API}/does-not-exist",
        json=INITIALIZE_MESSAGE,
        headers={"accept": CONTENT_TYPE_JSON},
    )
    assert response.status == HTTPStatus.NOT_FOUND
    assert "Unknown LLM API" in await response.text()


@pytest.mark.parametrize(
    ("require_admin", "expected_status"),
    [
        pytest.param(False, HTTPStatus.OK, id="not_required"),
        pytest.param(True, HTTPStatus.UNAUTHORIZED, id="required"),
    ],
)
async def test_require_admin_option(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
    hass_read_only_access_token: str,
    expected_status: HTTPStatus,
) -> None:
    """Test the require admin option applied to a non-admin user."""
    client = await hass_client(hass_read_only_access_token)

    response = await client.post(
        STREAMABLE_API,
        json=INITIALIZE_MESSAGE,
        headers={"accept": CONTENT_TYPE_JSON},
    )
    assert response.status == expected_status


@pytest.mark.parametrize("require_admin", [True])
async def test_require_admin_blocks_sse_endpoints(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
    hass_read_only_access_token: str,
) -> None:
    """Test the require admin option applied to the SSE endpoints."""
    client = await hass_client(hass_read_only_access_token)

    response = await client.get(SSE_API)
    assert response.status == HTTPStatus.UNAUTHORIZED

    response = await client.post(MESSAGES_API.format(session_id="session-id"))
    assert response.status == HTTPStatus.UNAUTHORIZED


@pytest.mark.parametrize("require_admin", [True])
async def test_require_admin_allows_admin(
    hass: HomeAssistant,
    setup_integration: None,
    hass_client: ClientSessionGenerator,
) -> None:
    """Test an admin user may use the endpoint that requires an admin."""
    client = await hass_client()

    response = await client.post(
        STREAMABLE_API,
        json=INITIALIZE_MESSAGE,
        headers={"accept": CONTENT_TYPE_JSON},
    )
    assert response.status == HTTPStatus.OK