test_oauth_flow.py

import asyncio
import os
from urllib.parse import parse_qs, urlparse

os.environ.setdefault("MCP_API_KEY", "test-key")
os.environ.setdefault("PUBLIC_BASE_URL", "http://testserver")

import pytest
from fastmcp import Client
from mcp.types import CallToolResult, ImageContent, TextContent
from starlette.testclient import TestClient
import app as bridge_app
from app import SESSIONS, app, _normalize_remote_url, _remote_client, _remote_headers, mcp_connect, mcp_tool_list


class FakeRemoteClient:
    def __init__(self, result):
        self.result = result
        self.calls = []

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return False

    async def call_tool_mcp(self, name, arguments):
        self.calls.append((name, arguments))
        return self.result

    async def list_tools(self):
        return []


def test_oauth_metadata_and_flow():
    with TestClient(app) as c:
        assert c.get("/").json() == {"ok": True}

        r = c.get("/.well-known/oauth-protected-resource")
        assert r.status_code == 200
        assert r.json()["authorization_servers"] == ["http://testserver"]

        r = c.get("/.well-known/oauth-authorization-server")
        assert r.status_code == 200
        assert "registration_endpoint" in r.json()

        r = c.post("/oauth/register", json={"redirect_uris": ["https://chatgpt.com/oauth/callback"]})
        assert r.status_code == 201
        client_id = r.json()["client_id"]

        params = {
            "response_type": "code",
            "client_id": client_id,
            "redirect_uri": "https://chatgpt.com/oauth/callback",
            "state": "abc",
            "resource": "http://testserver/sse",
            "code_challenge": "verifier",
            "code_challenge_method": "plain",
        }
        r = c.get("/oauth/authorize", params=params)
        assert r.status_code == 200

        r = c.post("/oauth/authorize", data={**params, "access_key": "test-key"}, follow_redirects=False)
        assert r.status_code in (302, 307)
        location = r.headers["location"]
        q = parse_qs(urlparse(location).query)
        code = q["code"][0]

        r = c.post("/oauth/token", data={
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": "https://chatgpt.com/oauth/callback",
            "client_id": client_id,
            "code_verifier": "verifier",
            "resource": "http://testserver/sse",
        })
        assert r.status_code == 200
        access_token = r.json()["access_token"]

        r = c.get("/sse")
        assert r.status_code == 401
        assert "WWW-Authenticate" in r.headers

        r = c.get("/sse", headers={"Authorization": f"Bearer {access_token}"})
        assert r.status_code != 401


def test_remote_client_preserves_headers_for_http_transport():
    client = _remote_client("https://example.test/sse", {"Authorization": "Bearer token"})

    assert client.transport.headers == {"Authorization": "Bearer token"}
    assert client.transport.url == "https://example.test/sse"


def test_normalize_remote_url_adds_https_for_bare_remote_host():
    assert _normalize_remote_url("patdev-mcp-bridge.hf.space/sse") == "https://patdev-mcp-bridge.hf.space/sse"


def test_normalize_remote_url_adds_http_for_localhost_target():
    assert _normalize_remote_url("localhost:4000/mcp") == "http://localhost:4000/mcp"


def test_remote_headers_adds_bearer_and_api_key():
    headers = _remote_headers({}, bearer_token="secret-token", x_api_key="secret-key")

    assert headers["Authorization"] == "Bearer secret-token"
    assert headers["x-api-key"] == "secret-key"


def test_mcp_connect_stores_normalized_url(monkeypatch):
    fake_client = FakeRemoteClient(None)
    monkeypatch.setattr("app._remote_client", lambda url, headers=None: fake_client)
    result = asyncio.run(mcp_connect.fn("localhost:4000/mcp"))
    session_id = result["session_id"]

    try:
        assert result["url"] == "http://localhost:4000/mcp"
        assert result["verified"] is True
        assert result["tool_count"] == 0
        assert SESSIONS[session_id]["url"] == "http://localhost:4000/mcp"
    finally:
        SESSIONS.pop(session_id, None)


def test_mcp_connect_merges_remote_auth_headers(monkeypatch):
    captured = {}
    fake_client = FakeRemoteClient(None)

    def fake_remote_client(url, headers=None):
        captured["url"] = url
        captured["headers"] = headers
        return fake_client

    monkeypatch.setattr("app._remote_client", fake_remote_client)
    result = asyncio.run(mcp_connect.fn("https://example.test/sse", bearer_token="token-123", x_api_key="key-456"))
    session_id = result["session_id"]

    try:
        assert captured["url"] == "https://example.test/sse"
        assert captured["headers"]["Authorization"] == "Bearer token-123"
        assert captured["headers"]["x-api-key"] == "key-456"
        assert SESSIONS[session_id]["headers"] == captured["headers"]
    finally:
        SESSIONS.pop(session_id, None)


def test_mcp_connect_reports_target_url_on_verify_failure(monkeypatch):
    class FailingClient:
        async def __aenter__(self):
            raise RuntimeError("401 Unauthorized")

        async def __aexit__(self, exc_type, exc, tb):
            return False

    monkeypatch.setattr("app._remote_client", lambda url, headers=None: FailingClient())

    with pytest.raises(ValueError, match="https://patdev-mcp-bridge.hf.space/sse"):
        asyncio.run(mcp_connect.fn("https://patdev-mcp-bridge.hf.space/sse", bearer_token="bad-token"))


def test_mcp_tool_call_emits_native_image_content_end_to_end(monkeypatch):
    result = CallToolResult(
        content=[
            TextContent(type="text", text="ok"),
            ImageContent(type="image", data="AAAA", mimeType="image/png"),
        ],
        isError=False,
    )
    fake_client = FakeRemoteClient(result)
    session_id = "test-session"
    SESSIONS[session_id] = {"url": "https://example.test/sse", "headers": {}}
    monkeypatch.setattr("app._remote_client", lambda url, headers=None: fake_client)

    async def run_test():
        client = Client(bridge_app.mcp)
        async with client:
            returned = await client.call_tool_mcp(
                "mcp_tool_call",
                {"session_id": session_id, "tool_name": "Screenshot", "arguments": {}},
            )
        assert returned.content[0].type == "image"
        assert returned.content[0].mimeType == "image/png"
        assert len(returned.content) == 1
        assert returned.isError is False

    try:
        asyncio.run(run_test())
    finally:
        SESSIONS.pop(session_id, None)

    assert fake_client.calls == [("Screenshot", {})]


def test_mcp_tool_list_strips_remote_output_schema(monkeypatch):
    class FakeTool:
        def model_dump(self):
            return {
                "name": "mcp_windows_Screenshot",
                "description": "fake screenshot tool",
                "inputSchema": {"type": "object"},
                "outputSchema": {"type": "object", "properties": {"image": {"type": "string"}}},
            }

    class ListingClient:
        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb):
            return False

        async def list_tools(self):
            return [FakeTool()]

    session_id = "tool-list-session"
    SESSIONS[session_id] = {"url": "https://example.test/sse", "headers": {}}
    monkeypatch.setattr("app._remote_client", lambda url, headers=None: ListingClient())

    async def run_test():
        returned = await mcp_tool_list.fn(session_id)
        assert returned["tools"][0]["name"] == "mcp_windows_Screenshot"
        assert returned["tools"][0]["inputSchema"] == {"type": "object"}
        assert "outputSchema" not in returned["tools"][0]

    try:
        asyncio.run(run_test())
    finally:
        SESSIONS.pop(session_id, None)


def test_mcp_tool_list_reports_target_url_on_connect_failure(monkeypatch):
    session_id = "broken-session"
    SESSIONS[session_id] = {"url": "https://missing.example.test/sse", "headers": {}}

    class FailingClient:
        async def __aenter__(self):
            raise OSError("[Errno -2] Name or service not known")

        async def __aexit__(self, exc_type, exc, tb):
            return False

    monkeypatch.setattr("app._remote_client", lambda url, headers=None: FailingClient())

    async def run_test():
        try:
            await mcp_tool_list.fn(session_id)
        except ValueError as exc:
            message = str(exc)
            assert "https://missing.example.test/sse" in message
            assert "Name or service not known" in message
        else:
            raise AssertionError("Expected mcp_tool_list to raise ValueError")

    try:
        asyncio.run(run_test())
    finally:
        SESSIONS.pop(session_id, None)


if __name__ == "__main__":
    test_oauth_metadata_and_flow()
    test_remote_client_preserves_headers_for_http_transport()
    print("OAuth smoke test passed")