tests/test_http_middleware.py

# Copyright 2026 Dataiku SAS
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""The HTTP middleware exchanges a DSS token only for DSS-backed tools."""

import asyncio
from types import SimpleNamespace

import pytest

import dataiku_mcp.server as server
from dataiku_mcp.config import request
from dataiku_mcp.instance_policy import (
    INSTANCE_MANAGEMENT_POLICY,
    PROJECT_READ_POLICY,
    PROJECT_EDIT_POLICY,
)
from fastmcp.exceptions import NotFoundError


def _call_middleware(
    monkeypatch,
    tool_name: str,
    *,
    tool_tags: set[str] | None = None,
    tool_resolves: bool = True,
    include_fastmcp_context: bool = True,
    instance_type: str = "design",
    tool_meta: dict | None = None,
    expected_error: tuple[type[Exception], str] | None = None,
):
    events = []
    access_token = SimpleNamespace(
        token="mcp-token",
        claims={"iss": "https://idp.example", "sub": "alice"},
    )
    middleware = server.RequestContextMiddleware()

    monkeypatch.setattr(server, "get_access_token", lambda: access_token)
    monkeypatch.setattr(
        request,
        "bind_http_identity",
        lambda issuer, subject: (
            events.append(("identity", issuer, subject)) or "identity"
        ),
    )
    monkeypatch.setattr(
        request,
        "pin_current_instance",
        lambda: events.append(("pin",)) or "instance",
    )
    monkeypatch.setattr(
        request,
        "reset_pinned_instance",
        lambda token: events.append(("reset_instance", token)),
    )
    monkeypatch.setattr(
        request,
        "reset_http_identity",
        lambda token: events.append(("reset_identity", token)),
    )
    monkeypatch.setattr(
        server,
        "exchange_http_token",
        lambda token: _exchange(events, token),
    )
    monkeypatch.setattr(
        request,
        "bind_http_dss_token",
        lambda token: events.append(("set_dss", token)) or "dss",
    )
    monkeypatch.setattr(
        request,
        "reset_http_dss_token",
        lambda token: events.append(("reset_dss", token)),
    )

    monkeypatch.setattr(
        request,
        "get_pinned_instance",
        lambda: SimpleNamespace(instance_type=instance_type),
    )

    async def call_next(_context):
        events.append(("tool",))
        return "result"

    class FakeFastMCP:
        async def get_tool(self, name):
            assert name == tool_name
            if not tool_resolves:
                return None
            return SimpleNamespace(
                tags=tool_tags or set(),
                meta=tool_meta
                or (
                    INSTANCE_MANAGEMENT_POLICY
                    if tool_name in {"switch_instance", "get_current_instance"}
                    else PROJECT_READ_POLICY
                ),
            )

    fastmcp_context = None
    if include_fastmcp_context:
        fastmcp_context = SimpleNamespace(fastmcp=FakeFastMCP())
    monkeypatch.setattr(server.mcp, "get_tool", FakeFastMCP().get_tool)
    context = SimpleNamespace(
        message=SimpleNamespace(name=tool_name),
        fastmcp_context=fastmcp_context,
    )
    if expected_error:
        with pytest.raises(expected_error[0], match=expected_error[1]):
            asyncio.run(middleware.on_call_tool(context, call_next))
    else:
        assert asyncio.run(middleware.on_call_tool(context, call_next)) == "result"
    return events


async def _exchange(events, token):
    events.append(("exchange", token))
    return "dss-token"


@pytest.mark.parametrize(
    "claims",
    [
        {"sub": "alice"},
        {"iss": "https://idp.example"},
        {"iss": None, "sub": "alice"},
        {"iss": 42, "sub": "alice"},
        {"iss": [], "sub": "alice"},
        {"iss": "https://idp.example", "sub": None},
        {"iss": "https://idp.example", "sub": 42},
        {"iss": "https://idp.example", "sub": []},
    ],
)
def test_http_invalid_identity_claims_fail_before_request_setup(monkeypatch, claims):
    events = []
    middleware = server.RequestContextMiddleware()
    access_token = SimpleNamespace(token="mcp-token", claims=claims)

    monkeypatch.setattr(server, "get_access_token", lambda: access_token)
    monkeypatch.setattr(
        request, "pin_current_instance", lambda: events.append(("pin",))
    )
    monkeypatch.setattr(
        server, "exchange_http_token", lambda _: events.append(("exchange",))
    )

    async def call_next(_context):
        events.append(("tool",))

    context = SimpleNamespace(message=SimpleNamespace(name="list_projects"))
    with pytest.raises(ValueError, match="non-empty string iss and sub claims"):
        asyncio.run(middleware.on_call_tool(context, call_next))

    assert events == []


def test_dss_independent_tools_are_tagged():
    async def get_tagged_tool_names():
        tools = await server.mcp.list_tools(run_middleware=False)
        return {
            tool.name for tool in tools if server.DSS_INDEPENDENT_TOOL_TAG in tool.tags
        }

    assert asyncio.run(get_tagged_tool_names()) == {
        "list_instances",
        "switch_instance",
        "delete_instance",
        "configure_instance",
        "get_cobuild_turn_status",
        "list_cobuild_conversations",
    }


def test_http_local_only_tool_skips_dss_token_exchange(monkeypatch):
    events = _call_middleware(
        monkeypatch,
        "switch_instance",
        tool_tags={server.DSS_INDEPENDENT_TOOL_TAG},
    )

    assert ("exchange", "mcp-token") not in events
    assert ("set_dss", "dss-token") not in events
    assert events[-2:] == [
        ("reset_instance", "instance"),
        ("reset_identity", "identity"),
    ]


def test_http_dss_tool_exchanges_and_resets_its_dss_token(monkeypatch):
    events = _call_middleware(monkeypatch, "list_projects")

    assert ("exchange", "mcp-token") in events
    assert ("set_dss", "dss-token") in events
    assert events[-3:] == [
        ("reset_dss", "dss"),
        ("reset_instance", "instance"),
        ("reset_identity", "identity"),
    ]


def test_http_current_instance_exchanges_a_dss_token(monkeypatch):
    events = _call_middleware(monkeypatch, "get_current_instance")

    assert ("exchange", "mcp-token") in events
    assert ("set_dss", "dss-token") in events


@pytest.mark.parametrize("include_fastmcp_context", [True, False])
def test_http_unknown_tool_fails_without_token_exchange(
    monkeypatch, include_fastmcp_context
):
    events = _call_middleware(
        monkeypatch,
        "unknown_tool",
        tool_resolves=False,
        include_fastmcp_context=include_fastmcp_context,
        expected_error=(NotFoundError, "Unknown tool"),
    )
    assert ("exchange", "mcp-token") not in events
    assert ("tool",) not in events
    assert events[-2:] == [
        ("reset_instance", "instance"),
        ("reset_identity", "identity"),
    ]


@pytest.mark.parametrize("instance_type", ["govern", "deployer", "agent-management"])
def test_http_unsupported_current_instance_skips_exchange(monkeypatch, instance_type):
    events = _call_middleware(
        monkeypatch, "get_current_instance", instance_type=instance_type
    )
    assert ("exchange", "mcp-token") not in events
    assert ("tool",) in events


@pytest.mark.parametrize(
    "tool_name,tool_tags",
    [
        ("set_project_variables", set()),
        ("list_cobuild_conversations", {server.DSS_INDEPENDENT_TOOL_TAG}),
    ],
)
def test_http_policy_rejection_precedes_exchange_and_handler(
    monkeypatch, tool_name, tool_tags
):
    events = _call_middleware(
        monkeypatch,
        tool_name,
        instance_type="automation",
        tool_tags=tool_tags,
        tool_meta=PROJECT_EDIT_POLICY,
        expected_error=(ValueError, "automation.*Allowed types.*switch_instance"),
    )
    assert ("exchange", "mcp-token") not in events
    assert ("tool",) not in events
    assert events[-2:] == [
        ("reset_instance", "instance"),
        ("reset_identity", "identity"),
    ]


def test_http_identity_is_reset_when_instance_pinning_fails(monkeypatch):
    events = []
    access_token = SimpleNamespace(
        token="mcp-token",
        claims={"iss": "https://idp.example", "sub": "alice"},
    )
    middleware = server.RequestContextMiddleware()

    monkeypatch.setattr(server, "get_access_token", lambda: access_token)
    monkeypatch.setattr(request, "bind_http_identity", lambda *_: "identity")

    def fail_pinning():
        raise ValueError("invalid settings")

    monkeypatch.setattr(
        request,
        "pin_current_instance",
        fail_pinning,
    )
    monkeypatch.setattr(
        request,
        "reset_http_identity",
        lambda token: events.append(("reset_identity", token)),
    )

    async def call_next(_context):
        pytest.fail("The tool must not run after pinning fails.")

    context = SimpleNamespace(message=SimpleNamespace(name="list_projects"))
    with pytest.raises(ValueError, match="invalid settings"):
        asyncio.run(middleware.on_call_tool(context, call_next))

    assert events == [("reset_identity", "identity")]