xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-09-16 23:07:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mcp): detect OAuth-protected remote MCP servers (#2363)

Browse Source 
* feat(mcp): surface OAuth-required server tests

* fix(mcp): show connection failure details in status cards

---------

Co-authored-by: RockChinQ <rockchinq@gmail.com>
This commit is contained in:
huige66631
2026-09-09 16:13:33 +08:00
committed by GitHub
parent 485113ae43
commit ce6b647fe7
18 changed files with 433 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit_tests/api/service/test_mcp_service.py
+70
View File
@@ -1009,6 +1009,37 @@ class TestMCPServiceTestMCPServer:
# Verify - returns task ID
assert task_id == 123
@pytest.mark.parametrize('refresh_first', [False, True])
async def test_persisted_test_preserves_failure_details(self, refresh_first):
from langbot.pkg.provider.tools.loaders.mcp import MCPSessionStatus
runtime_info = {'status': 'error', 'error_message': 'HTTP 403: access denied'}
session = SimpleNamespace(
status=MCPSessionStatus.CONNECTED if refresh_first else MCPSessionStatus.ERROR,
session=object(),
refresh=AsyncMock(side_effect=RuntimeError('refresh failed')),
start=AsyncMock(side_effect=RuntimeError('Connection failed, please check URL')),
get_runtime_info_dict=Mock(return_value=runtime_info),
)
captured = {}
def create_user_task(coroutine, **kwargs):
captured.update(coroutine=coroutine, context=kwargs['context'])
return SimpleNamespace(id=123)
ap = SimpleNamespace(
tool_mgr=SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=session))),
task_mgr=SimpleNamespace(create_user_task=Mock(side_effect=create_user_task)),
)
service = _service(ap)
service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'existing-server'}))
await service.test_mcp_server(_CONTEXT, 'existing-server', {})
with pytest.raises(RuntimeError, match='Connection failed'):
await captured['coroutine']
assert captured['context'].metadata['runtime_info'] == runtime_info
session.start.assert_awaited_once()
assert session.refresh.await_count == int(refresh_first)
async def test_test_mcp_server_not_found_raises(self):
"""Raises ValueError when server not found."""
# Setup
@@ -1052,6 +1083,45 @@ class TestMCPServiceTestMCPServer:
ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_called_once()
assert task_id == 456
async def test_transient_test_preserves_runtime_info_after_connection_failure(self):
runtime_info = {
'status': 'error',
'error_phase': 'oauth_required',
'retry_count': 1,
}
mock_session = SimpleNamespace(
server_name='oauth-server',
start=AsyncMock(side_effect=RuntimeError('connection failed')),
get_runtime_info_dict=Mock(return_value=runtime_info),
shutdown=AsyncMock(),
)
ap = SimpleNamespace(
tool_mgr=SimpleNamespace(
mcp_tool_loader=SimpleNamespace(load_mcp_server=AsyncMock(return_value=mock_session))
)
)
captured: dict = {}
def create_user_task(coroutine, **kwargs):
captured['coroutine'] = coroutine
captured['context'] = kwargs['context']
return SimpleNamespace(id=457)
ap.task_mgr = SimpleNamespace(create_user_task=Mock(side_effect=create_user_task))
service = _service(ap)
task_id = await service.test_mcp_server(
_CONTEXT,
'_',
{'name': 'OAuth server', 'mode': 'remote', 'enable': True, 'extra_args': {}},
)
assert task_id == 457
with pytest.raises(RuntimeError, match='connection failed'):
await captured['coroutine']
assert captured['context'].metadata['runtime_info'] == runtime_info
mock_session.shutdown.assert_awaited_once_with()
async def test_rejected_transient_test_session_is_shut_down(self):
ap = SimpleNamespace()
mock_session = MagicMock()
tests/unit_tests/provider/test_mcp_remote_transport.py
+75 -5
View File
@@ -13,7 +13,15 @@ from aiohttp import web
from mcp import types as mcp_types
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession
from langbot.pkg.provider.tools.loaders.mcp import MCPSessionStatus, MCPToolCallTimeoutError, RuntimeMCPSession
from langbot.pkg.provider.tools.loaders.mcp_stdio import MCPSessionErrorPhase
TEST_EXECUTION_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
TEST_EXECUTION_CONTEXT = ExecutionContext(
@@ -24,8 +32,9 @@ TEST_EXECUTION_CONTEXT = ExecutionContext(
class _TransportProbe:
def __init__(self, streamable_status: int | None) -> None:
def __init__(self, streamable_status: int | None, streamable_headers: dict[str, str] | None = None) -> None:
self.streamable_status = streamable_status
self.streamable_headers = streamable_headers or {}
self.streamable_posts = 0
self.streamable_messages: list[str] = []
self.sse_gets = 0
@@ -93,7 +102,7 @@ class _TransportProbe:
}
)
return web.Response(status=202)
return web.Response(status=self.streamable_status)
return web.Response(status=self.streamable_status, headers=self.streamable_headers)
self.sse_gets += 1
response = web.StreamResponse(
@@ -136,8 +145,8 @@ class _TransportProbe:
@asynccontextmanager
async def _transport_server(streamable_status: int | None):
probe = _TransportProbe(streamable_status)
async def _transport_server(streamable_status: int | None, streamable_headers: dict[str, str] | None = None):
probe = _TransportProbe(streamable_status, streamable_headers)
application = web.Application()
application.router.add_route('*', '/mcp', probe.handle_mcp_endpoint)
application.router.add_post('/messages', probe.handle_sse_message)
@@ -265,6 +274,45 @@ async def test_remote_transport_real_non_compatibility_error_does_not_fallback(s
await _close_session(session)
def test_remote_transport_extracts_oauth_resource_metadata_from_bearer_challenge():
request = httpx.Request('POST', 'https://mcp.example/mcp')
response = httpx.Response(
401,
headers={
'WWW-Authenticate': (
'Basic realm="MCP", Bearer resource_metadata="https://mcp.example/.well-known/oauth-protected-resource"'
)
},
request=request,
)
with pytest.raises(httpx.HTTPStatusError) as exc_info:
response.raise_for_status()
challenge = RuntimeMCPSession._extract_oauth_challenge(exc_info.value)
assert challenge is not None
assert challenge.resource_metadata_url == 'https://mcp.example/.well-known/oauth-protected-resource'
@pytest.mark.asyncio
async def test_remote_transport_oauth_challenge_sets_non_retryable_authorization_state():
headers = {
'WWW-Authenticate': 'Bearer resource_metadata="https://mcp.example/.well-known/oauth-protected-resource"'
}
async with _transport_server(401, headers) as (probe, url):
session = _session(url)
await session._lifecycle_loop_with_retry()
assert session.status == MCPSessionStatus.ERROR
assert session.error_phase == MCPSessionErrorPhase.OAUTH_REQUIRED
assert session.retry_count == 1
assert session._ready_event.is_set()
assert probe.streamable_posts == 1
assert probe.sse_gets == 0
@pytest.mark.asyncio
async def test_remote_transport_real_timeout_does_not_fallback():
async with _transport_server(None) as (probe, url):
@@ -313,3 +361,25 @@ async def test_remote_transport_external_cancellation_is_not_converted_to_sse_fa
finally:
probe.release_streamable_request.set()
await _close_session(session)
@pytest.mark.parametrize(
('error', 'expected'),
[
(httpx.ConnectError('secret host'), 'connection_unreachable'),
(httpx.ReadTimeout('secret URL'), 'connection_timeout'),
(TimeoutError('secret command'), 'connection_timeout'),
(RuntimeError('secret environment'), 'runtime_error'),
(
httpx.HTTPStatusError(
'secret response',
request=httpx.Request('POST', 'https://example.test/?token=secret'),
response=httpx.Response(403),
),
'http_403',
),
],
)
def test_public_error_category_does_not_expose_exception_details(error, expected):
grouped = ExceptionGroup('secret outer exception', [error])
assert RuntimeMCPSession._classify_public_error(grouped) == expected