test_mcp_oauth_e2e.py

import asyncio
import json
import os
from pathlib import Path
from typing import Any

import httpx
import pytest

from copilot.generated.rpc import MCPAppsCallToolRequest, MCPListToolsRequest
from copilot.session import MCPServerConfig, PermissionHandler
from copilot.session_events import McpServerStatus

from .testharness import E2ETestContext, wait_for_condition

TEST_MCP_OAUTH_SERVER = str(
    (Path(__file__).parents[2] / "test" / "harness" / "test-mcp-oauth-server.mjs").resolve()
)
EXPECTED_TOKEN = "sdk-host-token"
REFRESH_TOKEN = f"{EXPECTED_TOKEN}-refresh"
UPSCOPE_TOKEN = f"{EXPECTED_TOKEN}-upscope"
REAUTH_TOKEN = f"{EXPECTED_TOKEN}-reauth"

pytestmark = pytest.mark.asyncio(loop_scope="module")


async def _start_oauth_mcp_server() -> tuple[str, asyncio.subprocess.Process]:
    process = await asyncio.create_subprocess_exec(
        "node",
        TEST_MCP_OAUTH_SERVER,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env={**os.environ, "EXPECTED_TOKEN": EXPECTED_TOKEN},
    )
    assert process.stdout is not None

    try:
        line = await asyncio.wait_for(process.stdout.readline(), timeout=10)
    except TimeoutError as exc:
        await _stop_process(process)
        assert process.stderr is not None
        stderr = (await process.stderr.read()).decode(errors="replace")
        raise TimeoutError(f"Timed out waiting for OAuth MCP server: {stderr}") from exc
    if not line:
        assert process.stderr is not None
        stderr = (await process.stderr.read()).decode(errors="replace")
        raise RuntimeError(f"OAuth MCP server exited before listening: {stderr}")
    text = line.decode().strip()
    if text.startswith("Listening: "):
        return text.removeprefix("Listening: "), process

    await _stop_process(process)
    raise RuntimeError(f"Unexpected OAuth MCP server startup line: {text}")


async def _stop_process(process: asyncio.subprocess.Process) -> None:
    if process.returncode is not None:
        return
    process.terminate()
    try:
        await asyncio.wait_for(process.wait(), timeout=5)
    except TimeoutError:
        process.kill()
        await process.wait()


async def _requests(base_url: str) -> list[dict[str, Any]]:
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{base_url}/__requests")
        response.raise_for_status()
        return response.json()


async def _wait_for_mcp_server_status(
    session, server_name: str, expected_status: McpServerStatus = McpServerStatus.CONNECTED
) -> None:
    last_status = "<not listed>"

    async def matches() -> bool:
        nonlocal last_status
        result = await session.rpc.mcp.list()
        server = next((s for s in result.servers if s.name == server_name), None)
        last_status = server.status.value if server is not None else "<not listed>"
        return server is not None and server.status == expected_status

    await wait_for_condition(
        matches,
        timeout=60.0,
        poll_interval=0.2,
        timeout_message=(
            f"{server_name} did not reach {expected_status.value}; last status was {last_status}"
        ),
    )


class TestMcpOAuth:
    async def test_should_satisfy_mcp_oauth_using_host_provided_token(self, ctx: E2ETestContext):
        url, process = await _start_oauth_mcp_server()
        server_name = "oauth-protected-mcp"
        observed_request = None

        def on_mcp_auth_request(request, _invocation):
            nonlocal observed_request
            observed_request = request
            return {
                "kind": "token",
                "accessToken": EXPECTED_TOKEN,
                "tokenType": "Bearer",
                "expiresIn": 3600,
            }

        try:
            mcp_servers: dict[str, MCPServerConfig] = {
                server_name: {
                    "type": "http",
                    "url": f"{url}/mcp",
                    "tools": ["*"],
                }
            }
            async with await ctx.client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                on_mcp_auth_request=on_mcp_auth_request,
                mcp_servers=mcp_servers,
            ) as session:
                await _wait_for_mcp_server_status(session, server_name)

                tools = await session.rpc.mcp.list_tools(
                    MCPListToolsRequest(server_name=server_name)
                )
                assert [tool.name for tool in tools.tools] == ["whoami"]

            assert observed_request is not None
            assert observed_request["serverName"] == server_name
            assert observed_request["serverUrl"] == f"{url}/mcp"
            assert observed_request["reason"] == "initial"
            assert observed_request["wwwAuthenticateParams"] == {
                "resourceMetadataUrl": f"{url}/.well-known/oauth-protected-resource",
                "scope": "mcp.read",
                "error": "invalid_token",
            }
            assert json.loads(observed_request["resourceMetadata"]) == {
                "resource": f"{url}/mcp",
                "authorization_servers": [url],
                "scopes_supported": ["mcp.read"],
                "bearer_methods_supported": ["header"],
            }

            requests = await _requests(url)
            assert any(request["authorization"] is None for request in requests)
            assert any(
                request["authorization"] == f"Bearer {EXPECTED_TOKEN}" for request in requests
            )
        finally:
            await _stop_process(process)

    async def test_should_request_replacement_tokens_across_mcp_oauth_lifecycle(
        self, ctx: E2ETestContext
    ):
        url, process = await _start_oauth_mcp_server()
        server_name = "oauth-lifecycle-mcp"
        observed_requests: list[dict[str, Any]] = []
        refresh_count = 0

        def on_mcp_auth_request(request, _invocation):
            nonlocal refresh_count
            observed_requests.append(request)
            if request["reason"] == "refresh":
                refresh_count += 1
                assert request["wwwAuthenticateParams"] == {"error": "invalid_token"}
                if refresh_count > 1:
                    return {"kind": "cancelled"}
                return {"kind": "token", "accessToken": REFRESH_TOKEN}
            if request["reason"] == "upscope":
                assert request["wwwAuthenticateParams"] == {
                    "resourceMetadataUrl": f"{url}/.well-known/oauth-protected-resource",
                    "scope": "mcp.write",
                    "error": "insufficient_scope",
                }
                return {"kind": "token", "accessToken": UPSCOPE_TOKEN}
            if request["reason"] == "reauth":
                return {"kind": "token", "accessToken": REAUTH_TOKEN}
            return {"kind": "token", "accessToken": EXPECTED_TOKEN}

        try:
            mcp_servers: dict[str, MCPServerConfig] = {
                server_name: {
                    "type": "http",
                    "url": f"{url}/mcp",
                    "tools": ["*"],
                }
            }
            async with await ctx.client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                on_mcp_auth_request=on_mcp_auth_request,
                mcp_servers=mcp_servers,
                enable_mcp_apps=True,
            ) as session:
                await _wait_for_mcp_server_status(session, server_name)

                for scenario in ("refresh", "upscope", "reauth"):
                    result = await session.rpc.mcp.apps.call_tool(
                        MCPAppsCallToolRequest(
                            origin_server_name=server_name,
                            server_name=server_name,
                            tool_name="whoami",
                            arguments={"scenario": scenario},
                        )
                    )
                    assert result["content"] == [{"type": "text", "text": "oauth-test-user"}]

            assert [request["reason"] for request in observed_requests] == [
                "initial",
                "refresh",
                "upscope",
                "refresh",
                "reauth",
            ]
            requests = await _requests(url)
            assert any(
                request["authorization"] == f"Bearer {REFRESH_TOKEN}" for request in requests
            )
            assert any(
                request["authorization"] == f"Bearer {UPSCOPE_TOKEN}" for request in requests
            )
            assert any(request["authorization"] == f"Bearer {REAUTH_TOKEN}" for request in requests)
        finally:
            await _stop_process(process)

    async def test_should_cancel_pending_mcp_oauth_request(self, ctx: E2ETestContext):
        url, process = await _start_oauth_mcp_server()
        server_name = "oauth-cancelled-mcp"
        observed_request = None

        def on_mcp_auth_request(request, _invocation):
            nonlocal observed_request
            observed_request = request
            return {"kind": "cancelled"}

        try:
            mcp_servers: dict[str, MCPServerConfig] = {
                server_name: {
                    "type": "http",
                    "url": f"{url}/mcp",
                    "tools": ["*"],
                }
            }
            async with await ctx.client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                on_mcp_auth_request=on_mcp_auth_request,
                mcp_servers=mcp_servers,
            ) as session:
                await _wait_for_mcp_server_status(session, server_name, McpServerStatus.FAILED)

            assert observed_request is not None
            assert observed_request["serverName"] == server_name
            assert observed_request["reason"] == "initial"
        finally:
            await _stop_process(process)