mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 00:03:01 +03:00
chengyongru
be3a42ebac
Move blocking filesystem, persistence, subprocess, media, and DNS work off the gateway event loop while preserving existing contracts. Add bounded cancellation and responsiveness regression coverage.
298 lines
9.7 KiB
Python
298 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from nanobot.config.schema import MCPServerConfig
|
|
from nanobot.webui.mcp_oauth_api import (
|
|
McpOAuthError,
|
|
McpOAuthManager,
|
|
prepare_mcp_oauth_redirect_uri,
|
|
validate_mcp_oauth_redirect_uri,
|
|
)
|
|
|
|
|
|
class _Connection:
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
|
|
async def aclose(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
def _config() -> MCPServerConfig:
|
|
return MCPServerConfig(
|
|
type="streamableHttp",
|
|
auth="oauth",
|
|
url="https://mcp.example.com/mcp",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_browser_flow_retries_current_server_and_ignores_unrelated_reload_failure(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
manager = McpOAuthManager()
|
|
connection = _Connection()
|
|
received: dict[str, object] = {}
|
|
reload_calls = 0
|
|
|
|
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
|
return True, ""
|
|
|
|
monkeypatch.setattr(
|
|
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
|
validate_url_target,
|
|
)
|
|
|
|
async def connect(servers, _registry, *, oauth_handlers):
|
|
assert set(servers) == {"xmind"}
|
|
handlers = oauth_handlers["xmind"]
|
|
await handlers.redirect_handler(
|
|
"https://accounts.example.com/authorize?client_id=test&state=state-123"
|
|
)
|
|
received["callback"] = await handlers.callback_handler()
|
|
return {"xmind": connection}
|
|
|
|
async def reload_mcp() -> dict[str, object]:
|
|
nonlocal reload_calls
|
|
reload_calls += 1
|
|
if reload_calls == 1:
|
|
return {
|
|
"ok": False,
|
|
"requires_restart": False,
|
|
"failed": ["xmind"],
|
|
}
|
|
return {
|
|
"ok": False,
|
|
"requires_restart": False,
|
|
"connected": ["xmind"],
|
|
"failed": ["notion"],
|
|
"message": "MCP config reloaded, but some servers did not connect: notion",
|
|
}
|
|
|
|
monkeypatch.setattr("nanobot.webui.mcp_oauth_api.connect_mcp_servers", connect)
|
|
|
|
started = await manager.start(
|
|
"xmind",
|
|
_config(),
|
|
"https://agent.example.com/auth/mcp/callback",
|
|
reload_mcp=reload_mcp,
|
|
)
|
|
|
|
assert started["status"] == "authorization_required"
|
|
assert started["authorization_url"].startswith("https://accounts.example.com/authorize?")
|
|
manager.submit_callback(state="state-123", code="oauth-code", error=None)
|
|
with pytest.raises(McpOAuthError, match="expired"):
|
|
manager.submit_callback(state="state-123", code="replayed-code", error=None)
|
|
|
|
for _ in range(10):
|
|
await asyncio.sleep(0)
|
|
if reload_calls == 2:
|
|
break
|
|
assert reload_calls == 2
|
|
|
|
for _ in range(10):
|
|
first, second = await asyncio.gather(
|
|
manager.status(started["flow_id"]),
|
|
manager.status(started["flow_id"]),
|
|
)
|
|
if first["status"] == "connected":
|
|
break
|
|
await asyncio.sleep(0)
|
|
|
|
assert first["status"] == "connected"
|
|
assert second["status"] == "connected"
|
|
assert first["hot_reload"]["failed"] == ["notion"]
|
|
assert received["callback"] == ("oauth-code", "state-123")
|
|
assert reload_calls == 2
|
|
assert connection.closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_remote_http_flow_accepts_a_pasted_loopback_callback(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
manager = McpOAuthManager()
|
|
connection = _Connection()
|
|
received: dict[str, object] = {}
|
|
|
|
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
|
return True, ""
|
|
|
|
monkeypatch.setattr(
|
|
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
|
validate_url_target,
|
|
)
|
|
|
|
async def connect(_servers, _registry, *, oauth_handlers):
|
|
handlers = oauth_handlers["linear"]
|
|
received["redirect_uri"] = handlers.redirect_uri
|
|
await handlers.redirect_handler(
|
|
"https://accounts.example.com/authorize?client_id=test&state=manual-state"
|
|
)
|
|
received["callback"] = await handlers.callback_handler()
|
|
return {"linear": connection}
|
|
|
|
monkeypatch.setattr("nanobot.webui.mcp_oauth_api.connect_mcp_servers", connect)
|
|
|
|
started = await manager.start(
|
|
"linear",
|
|
_config(),
|
|
"http://192.0.2.10:8765/auth/mcp/callback",
|
|
reload_mcp=lambda: asyncio.sleep(
|
|
0,
|
|
result={"ok": True, "requires_restart": False},
|
|
),
|
|
)
|
|
|
|
assert started["status"] == "authorization_required"
|
|
assert started["completion_input"] == "callback_url"
|
|
assert received["redirect_uri"] == "http://127.0.0.1:8765/auth/mcp/callback"
|
|
|
|
with pytest.raises(McpOAuthError, match="complete callback URL"):
|
|
manager.submit_callback_url(
|
|
flow_id=started["flow_id"],
|
|
callback_url=(
|
|
"http://127.0.0.1:8765/wrong?code=oauth-code&state=manual-state"
|
|
),
|
|
)
|
|
with pytest.raises(McpOAuthError, match="different or expired"):
|
|
manager.submit_callback_url(
|
|
flow_id=started["flow_id"],
|
|
callback_url=(
|
|
"http://127.0.0.1:8765/auth/mcp/callback"
|
|
"?code=oauth-code&state=other-state"
|
|
),
|
|
)
|
|
|
|
submitted = manager.submit_callback_url(
|
|
flow_id=started["flow_id"],
|
|
callback_url=(
|
|
"http://127.0.0.1:8765/auth/mcp/callback"
|
|
"?code=oauth-code&state=manual-state"
|
|
),
|
|
)
|
|
assert submitted["status"] == "connecting"
|
|
|
|
for _ in range(20):
|
|
await asyncio.sleep(0)
|
|
result = await manager.status(started["flow_id"])
|
|
if result["status"] == "connected":
|
|
break
|
|
|
|
assert result["status"] == "connected"
|
|
assert result["completion_input"] == "callback_url"
|
|
assert received["callback"] == ("oauth-code", "manual-state")
|
|
assert connection.closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_browser_flow_surfaces_provider_denial_without_callback_description(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
manager = McpOAuthManager()
|
|
|
|
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
|
return True, ""
|
|
|
|
monkeypatch.setattr(
|
|
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
|
validate_url_target,
|
|
)
|
|
|
|
async def connect(_servers, _registry, *, oauth_handlers):
|
|
handlers = oauth_handlers["notion"]
|
|
await handlers.redirect_handler("https://accounts.example.com/auth?state=deny-state")
|
|
await handlers.callback_handler()
|
|
return {}
|
|
|
|
monkeypatch.setattr("nanobot.webui.mcp_oauth_api.connect_mcp_servers", connect)
|
|
started = await manager.start(
|
|
"notion",
|
|
_config(),
|
|
"https://agent.example.com/auth/mcp/callback",
|
|
reload_mcp=lambda: asyncio.sleep(0, result={"ok": True}),
|
|
)
|
|
|
|
with pytest.raises(McpOAuthError, match="access_denied"):
|
|
manager.submit_callback(state="deny-state", code=None, error="access_denied")
|
|
for _ in range(10):
|
|
await asyncio.sleep(0)
|
|
result = await manager.status(started["flow_id"])
|
|
if result["status"] == "failed":
|
|
break
|
|
|
|
assert result["status"] == "failed"
|
|
assert result["error"] == "Authorization was not completed (access_denied)."
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("authorization_url", "url_is_safe", "state"),
|
|
[
|
|
("https://127.0.0.1/authorize?state=private-state", False, "private-state"),
|
|
("http://accounts.example.com/authorize?state=http-state", True, "http-state"),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_browser_flow_blocks_unsafe_authorization_url(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
authorization_url: str,
|
|
url_is_safe: bool,
|
|
state: str,
|
|
) -> None:
|
|
manager = McpOAuthManager()
|
|
|
|
async def validate_url_target(_url: str) -> tuple[bool, str]:
|
|
return url_is_safe, "private address"
|
|
|
|
monkeypatch.setattr(
|
|
"nanobot.webui.mcp_oauth_api.async_validate_url_target",
|
|
validate_url_target,
|
|
)
|
|
|
|
async def connect(_servers, _registry, *, oauth_handlers):
|
|
await oauth_handlers["linear"].redirect_handler(authorization_url)
|
|
return {}
|
|
|
|
monkeypatch.setattr("nanobot.webui.mcp_oauth_api.connect_mcp_servers", connect)
|
|
|
|
result = await manager.start(
|
|
"linear",
|
|
_config(),
|
|
"https://agent.example.com/auth/mcp/callback",
|
|
reload_mcp=lambda: asyncio.sleep(0, result={"ok": True}),
|
|
)
|
|
|
|
assert result["status"] == "failed"
|
|
assert result["error"] == "The MCP server returned an unsafe authorization URL."
|
|
with pytest.raises(McpOAuthError, match="expired"):
|
|
manager.submit_callback(state=state, code="code", error=None)
|
|
|
|
|
|
def test_redirect_uri_requires_https_except_for_loopback() -> None:
|
|
assert validate_mcp_oauth_redirect_uri(
|
|
"https://agent.example.com/auth/mcp/callback"
|
|
) == "https://agent.example.com/auth/mcp/callback"
|
|
assert validate_mcp_oauth_redirect_uri(
|
|
"http://127.0.0.1:8765/auth/mcp/callback"
|
|
) == "http://127.0.0.1:8765/auth/mcp/callback"
|
|
|
|
with pytest.raises(McpOAuthError, match="HTTPS or localhost"):
|
|
validate_mcp_oauth_redirect_uri("http://192.0.2.10/auth/mcp/callback")
|
|
with pytest.raises(McpOAuthError, match="Invalid"):
|
|
validate_mcp_oauth_redirect_uri("https://agent.example.com/wrong")
|
|
|
|
|
|
def test_remote_http_redirect_prepares_a_manual_loopback_callback() -> None:
|
|
assert prepare_mcp_oauth_redirect_uri(
|
|
"https://agent.example.com/auth/mcp/callback"
|
|
) == ("https://agent.example.com/auth/mcp/callback", False)
|
|
assert prepare_mcp_oauth_redirect_uri(
|
|
"http://127.0.0.1:8765/auth/mcp/callback"
|
|
) == ("http://127.0.0.1:8765/auth/mcp/callback", False)
|
|
assert prepare_mcp_oauth_redirect_uri(
|
|
"http://agent.example.com:9443/auth/mcp/callback"
|
|
) == ("http://127.0.0.1:9443/auth/mcp/callback", True)
|