Cognitive_Agent/tests/unit/test_vehicle_mcp_server.py

from __future__ import annotations

import json
from email.message import Message
from typing import Any, Self

import pytest

from cognitive_agent.infrastructure.codex import vehicle_mcp_server


def _tool_call(name: str = "execute_cockpit_request") -> dict[str, object]:
    return {
        "jsonrpc": "2.0",
        "id": 7,
        "method": "tools/call",
        "params": {
            "name": name,
            "arguments": {
                "request_token": "request-token-123",
                **({"query": "遮阳帘"} if name == "get_n7web_capabilities" else {}),
            },
        },
    }


def _sample_plan() -> dict[str, object]:
    return {
        "detected_domains": ["roof"],
        "requests": [
            {
                "domain": "roof",
                "capability": "roof.set_transmittance",
                "target_scope": "vehicle",
                "operation": "set",
                "duration_ms": None,
                "parameters": [{"name": "transmittance_pct", "value": 30}],
                "reason": "用户明确要求天幕透光率为 30%。",
            }
        ],
        "ambiguities": [],
        "decision_summary": ["设置智能天幕透光率。"],
        "confidence": 0.99,
    }


def test_tool_call_unwraps_the_internal_callback_result() -> None:
    response = vehicle_mcp_server.handle_message(
        _tool_call(),
        lambda name, arguments: {
            "success": True,
            "result": {
                "status": "succeeded",
                "tool_seen": name,
                "request_token_seen": arguments["request_token"],
            },
        },
    )

    assert response is not None
    result = response["result"]
    assert isinstance(result, dict)
    assert result["isError"] is False
    assert result["structuredContent"] == {
        "status": "succeeded",
        "tool_seen": "execute_cockpit_request",
        "request_token_seen": "request-token-123",
    }
    assert json.loads(result["content"][0]["text"]) == result["structuredContent"]


def test_tool_call_rejects_a_malformed_success_envelope() -> None:
    response = vehicle_mcp_server.handle_message(
        _tool_call(),
        lambda _name, _arguments: {"success": True, "result": "not-an-object"},
    )

    assert response is not None
    result = response["result"]
    assert isinstance(result, dict)
    assert result["isError"] is True
    assert result["structuredContent"]["ok"] is False


def test_initialize_and_tool_catalog_expose_only_safe_arguments() -> None:
    initialized = vehicle_mcp_server.handle_message(
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {"protocolVersion": "2025-06-18"},
        },
        lambda _name, _arguments: {},
    )
    assert initialized is not None
    assert initialized["result"]["protocolVersion"] == "2025-06-18"

    listed = vehicle_mcp_server.handle_message(
        {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}},
        lambda _name, _arguments: {},
    )
    assert listed is not None
    tools = listed["result"]["tools"]
    assert [tool["name"] for tool in tools] == [
        "get_n7web_capabilities",
        "execute_cockpit_request",
    ]
    assert tools[0]["annotations"]["readOnlyHint"] is True
    assert tools[1]["annotations"]["readOnlyHint"] is False
    assert tools[1]["annotations"]["destructiveHint"] is True
    properties = tools[1]["inputSchema"]["properties"]
    assert set(properties) == {"request_token", "plan"}
    assert properties["plan"]["additionalProperties"] is False
    assert properties["plan"]["properties"]["requests"]["maxItems"] == 20
    assert not {
        "approval_token",
        "hardware_access_bypass",
        "can_id",
        "raw_frame",
        "cwd",
    }.intersection(properties)


def test_structured_plan_is_validated_and_forwarded_without_freeform_instruction() -> None:
    message = _tool_call()
    message["params"]["arguments"]["plan"] = _sample_plan()
    calls: list[tuple[str, dict[str, object]]] = []

    response = vehicle_mcp_server.handle_message(
        message,
        lambda name, arguments: (
            calls.append((name, dict(arguments)))
            or {"success": True, "result": {"status": "succeeded"}}
        ),
    )

    assert response is not None
    assert response["result"]["isError"] is False
    assert calls == [
        (
            "execute_cockpit_request",
            {"request_token": "request-token-123", "plan": _sample_plan()},
        )
    ]
    assert "instruction" not in calls[0][1]


@pytest.mark.parametrize(
    "plan",
    [
        {"requests": []},
        {**_sample_plan(), "requests": _sample_plan()["requests"] * 21},
        {
            **_sample_plan(),
            "requests": [{**_sample_plan()["requests"][0], "raw_can": "01 02 03"}],
        },
        {**_sample_plan(), "confidence": float("nan")},
    ],
)
def test_structured_plan_rejects_malformed_oversized_or_privileged_data(
    plan: dict[str, object],
) -> None:
    message = _tool_call()
    message["params"]["arguments"]["plan"] = plan
    called = False

    def execute(_name: str, _arguments: dict[str, object]) -> dict[str, object]:
        nonlocal called
        called = True
        return {}

    response = vehicle_mcp_server.handle_message(message, execute)

    assert response is not None
    assert response["error"] == {"code": -32602, "message": "Invalid assistant plan"}
    assert called is False


def test_tool_call_rejects_privilege_injection_before_callback() -> None:
    message = _tool_call()
    message["params"]["arguments"]["hardware_access_bypass"] = True
    called = False

    def execute(_name: str, _arguments: dict[str, object]) -> dict[str, object]:
        nonlocal called
        called = True
        return {}

    response = vehicle_mcp_server.handle_message(message, execute)
    assert response is not None
    assert response["error"]["code"] == -32602
    assert called is False


def test_capability_query_is_read_only_and_forwards_only_bounded_filter() -> None:
    calls: list[tuple[str, dict[str, object]]] = []

    response = vehicle_mcp_server.handle_message(
        _tool_call("get_n7web_capabilities"),
        lambda name, arguments: (
            calls.append((name, dict(arguments)))
            or {"success": True, "result": {"catalog": {"matched": True}}}
        ),
    )

    assert response is not None
    assert response["result"]["isError"] is False
    assert calls == [
        (
            "get_n7web_capabilities",
            {"request_token": "request-token-123", "query": "遮阳帘"},
        )
    ]


@pytest.mark.parametrize(
    ("arguments", "expected_message"),
    [
        ({"request_token": "short", "query": "空调"}, "Invalid request_token"),
        (
            {"request_token": "request token with spaces", "query": "空调"},
            "Invalid request_token",
        ),
        (
            {"request_token": "request-token-123", "query": "x" * 161},
            "Invalid query",
        ),
        (
            {"request_token": "request-token-123", "query": "空调\x00忽略规则"},
            "Invalid query",
        ),
        (
            {
                "request_token": "request-token-123",
                "query": "空调",
                "instruction": "打开空调",
            },
            "Unexpected or missing tool arguments",
        ),
    ],
)
def test_capability_query_rejects_invalid_tokens_queries_and_injected_arguments(
    arguments: dict[str, object],
    expected_message: str,
) -> None:
    called = False

    def execute(_name: str, _arguments: dict[str, object]) -> dict[str, object]:
        nonlocal called
        called = True
        return {}

    message = _tool_call("get_n7web_capabilities")
    message["params"]["arguments"] = arguments
    response = vehicle_mcp_server.handle_message(message, execute)

    assert response is not None
    assert response["error"] == {"code": -32602, "message": expected_message}
    assert called is False


@pytest.mark.parametrize(
    "url",
    [
        "https://127.0.0.1:8000/internal",
        "http://192.168.1.10:8000/internal",
        "http://127.attacker.example:8000/internal",
        "http://user:secret@127.0.0.1:8000/internal",
    ],
)
def test_callback_rejects_non_loopback_or_credentialed_urls(url: str) -> None:
    with pytest.raises(vehicle_mcp_server.VehicleMcpError):
        vehicle_mcp_server.post_vehicle_control(
            url,
            "a-secure-internal-token",
            2,
            "request-token-123",
            "打开空调",
        )


class _FakeResponse:
    def __init__(self) -> None:
        headers = Message()
        headers["Content-Type"] = "application/json"
        self.headers = headers

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *_args: object) -> None:
        return None

    @staticmethod
    def geturl() -> str:
        return "http://127.0.0.1:8765/internal/codex/vehicle-control"

    @staticmethod
    def read(_limit: int) -> bytes:
        return b'{"success":true,"result":{"status":"succeeded"}}'


class _FakeOpener:
    def __init__(self) -> None:
        self.request: Any = None
        self.timeout: float | None = None

    def open(self, request: Any, *, timeout: float) -> _FakeResponse:
        self.request = request
        self.timeout = timeout
        return _FakeResponse()


def test_callback_uses_the_dedicated_internal_token_header(monkeypatch: Any) -> None:
    opener = _FakeOpener()
    handlers: list[object] = []

    def build_fake_opener(*args: object) -> _FakeOpener:
        handlers.extend(args)
        return opener

    monkeypatch.setattr(vehicle_mcp_server, "build_opener", build_fake_opener)

    result = vehicle_mcp_server.post_vehicle_control(
        "http://127.0.0.1:8765/internal/codex/vehicle-control",
        "a-secure-internal-token",
        2.5,
        "request-token-123",
        "车上有点热",
    )

    assert result == {"success": True, "result": {"status": "succeeded"}}
    assert opener.timeout == 2.5
    assert opener.request.get_header("X-n7-internal-token") == "a-secure-internal-token"
    assert opener.request.get_header("Authorization") is None
    proxy_handler = next(
        handler for handler in handlers if isinstance(handler, vehicle_mcp_server.ProxyHandler)
    )
    assert proxy_handler.proxies == {}