70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import socket
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import pytest
|
|
import uvicorn
|
|
from fastmcp import Client, FastMCP
|
|
from fastmcp.client.transports import StreamableHttpTransport
|
|
from fastmcp.server.dependencies import get_http_headers
|
|
|
|
from beaver_gateway.backends.claude_sdk import ORIGIN_HEADER
|
|
from beaver_gateway.mcp.internal_app import build_internal_app
|
|
from beaver_gateway.mcp.types import McpServer
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
|
|
def _free_port() -> int:
|
|
with socket.socket() as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def _serving(app: Any, port: int) -> AsyncIterator[None]:
|
|
server = uvicorn.Server(
|
|
uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
|
|
)
|
|
task = asyncio.create_task(server.serve())
|
|
while not server.started:
|
|
await asyncio.sleep(0.02)
|
|
try:
|
|
yield
|
|
finally:
|
|
server.should_exit = True
|
|
await task
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_the_origin_header_reaches_the_upstream_mcp() -> None:
|
|
upstream = FastMCP("upstream")
|
|
|
|
@upstream.tool
|
|
def caller() -> str:
|
|
return get_http_headers().get(ORIGIN_HEADER.lower(), "")
|
|
|
|
up, gateway = _free_port(), _free_port()
|
|
app, urls, _ = build_internal_app(
|
|
[McpServer.http(name="t3", url=f"http://127.0.0.1:{up}/mcp")],
|
|
host="127.0.0.1",
|
|
port=gateway,
|
|
)
|
|
seen = []
|
|
async with (
|
|
_serving(upstream.http_app(transport="http", path="/mcp"), up),
|
|
_serving(app, gateway),
|
|
):
|
|
for conversation in ("conv-a", "conv-b"):
|
|
transport = StreamableHttpTransport(
|
|
urls["t3"], headers={ORIGIN_HEADER: conversation}
|
|
)
|
|
async with Client(transport) as client:
|
|
result = await client.call_tool("caller", {})
|
|
seen.append(result.data)
|
|
assert seen == ["conv-a", "conv-b"]
|