v2.17.0 (#822)
Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Hdandria <henri.dandria@mistral.ai> Co-authored-by: Ivana Dunisijevic <ivana.dunisijevic@mistral.ai> Co-authored-by: Jean Burellier <sheplu@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Mert Unsal <mert.unsal@mistral.ai> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Paul VEZIA <166131032+le-codeur-rapide@users.noreply.github.com> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: renovate-mistral[bot] <253709520+renovate-mistral[bot]@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin
GitHub
parent
564a14365e
commit
6bedf271ce
223 changed files with 10533 additions and 6947 deletions
323
tests/tools/test_mcp_registry_oauth.py
Normal file
323
tests/tools/test_mcp_registry_oauth.py
Normal file
|
|
@ -0,0 +1,323 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Iterator
|
||||
from typing import Any, ClassVar
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import keyring
|
||||
from keyring.backend import KeyringBackend
|
||||
import keyring.errors
|
||||
from mcp.client.auth import OAuthFlowError
|
||||
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
|
||||
import pytest
|
||||
|
||||
from vibe.core.auth.mcp_oauth import (
|
||||
Fingerprint,
|
||||
KeyringTokenStorage,
|
||||
MCPOAuthLoginFailed,
|
||||
)
|
||||
from vibe.core.config import MCPOAuth, MCPStreamableHttp
|
||||
from vibe.core.tools.base import BaseToolConfig, InvokeContext, ToolError
|
||||
from vibe.core.tools.mcp import AuthStatus, MCPRegistry, MCPToolResult, RemoteTool
|
||||
from vibe.core.tools.mcp.tools import (
|
||||
MCPHttpOAuthRuntime,
|
||||
_OpenArgs,
|
||||
create_mcp_http_proxy_tool_class,
|
||||
)
|
||||
|
||||
|
||||
class MemoryKeyring(KeyringBackend):
|
||||
priority: ClassVar[Any] = 100
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.store: dict[tuple[str, str], str] = {}
|
||||
|
||||
def get_password(self, service: str, username: str) -> str | None:
|
||||
return self.store.get((service, username))
|
||||
|
||||
def set_password(self, service: str, username: str, password: str) -> None:
|
||||
self.store[(service, username)] = password
|
||||
|
||||
def delete_password(self, service: str, username: str) -> None:
|
||||
if (service, username) not in self.store:
|
||||
raise keyring.errors.PasswordDeleteError(username)
|
||||
del self.store[(service, username)]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_keyring() -> Iterator[MemoryKeyring]:
|
||||
original = keyring.get_keyring()
|
||||
fake = MemoryKeyring()
|
||||
keyring.set_keyring(fake)
|
||||
try:
|
||||
yield fake
|
||||
finally:
|
||||
keyring.set_keyring(original)
|
||||
|
||||
|
||||
def _oauth_server(
|
||||
*,
|
||||
name: str = "linear",
|
||||
url: str = "https://mcp.example.com/mcp",
|
||||
scopes: list[str] | None = None,
|
||||
) -> MCPStreamableHttp:
|
||||
return MCPStreamableHttp(
|
||||
name=name,
|
||||
transport="streamable-http",
|
||||
url=url,
|
||||
auth=MCPOAuth(type="oauth", scopes=scopes if scopes is not None else ["read"]),
|
||||
)
|
||||
|
||||
|
||||
async def _save_valid_oauth_state(srv: MCPStreamableHttp) -> None:
|
||||
storage = KeyringTokenStorage(alias=srv.name)
|
||||
await storage.set_tokens(
|
||||
OAuthToken(
|
||||
access_token="ACCESS",
|
||||
token_type="Bearer",
|
||||
expires_in=3600,
|
||||
refresh_token="REFRESH",
|
||||
)
|
||||
)
|
||||
await Fingerprint.compute(srv).save(srv.name)
|
||||
|
||||
|
||||
class TestMCPRegistryOAuthDiscovery:
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_tokens_mark_needs_auth_without_caching(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
registry = MCPRegistry()
|
||||
|
||||
tools = await registry.get_tools_async([srv])
|
||||
|
||||
assert tools == {}
|
||||
assert registry.needs_auth == {"linear"}
|
||||
assert registry.status()["linear"] == AuthStatus.NEEDS_AUTH
|
||||
assert registry.count_loaded([srv]) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_removed_oauth_server_clears_auth_status(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
registry = MCPRegistry()
|
||||
|
||||
await registry.get_tools_async([srv])
|
||||
await registry.get_tools_async([])
|
||||
|
||||
assert registry.needs_auth == set()
|
||||
assert registry.status() == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_valid_tokens_register_tools_with_oauth_provider(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
await _save_valid_oauth_state(srv)
|
||||
registry = MCPRegistry()
|
||||
remote = RemoteTool(name="create_issue")
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http",
|
||||
new=AsyncMock(return_value=[remote]),
|
||||
) as list_tools:
|
||||
tools = await registry.get_tools_async([srv])
|
||||
|
||||
assert "linear_create_issue" in tools
|
||||
assert registry.needs_auth == set()
|
||||
assert registry.status()["linear"] == AuthStatus.OK
|
||||
await_args = list_tools.await_args
|
||||
assert await_args is not None
|
||||
assert await_args.kwargs["auth"] is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fingerprint_mismatch_deletes_tokens_and_requires_auth(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
original = _oauth_server(scopes=["read"])
|
||||
changed = _oauth_server(scopes=["read", "write"])
|
||||
await _save_valid_oauth_state(original)
|
||||
storage = KeyringTokenStorage(alias="linear")
|
||||
await storage.set_client_info(
|
||||
OAuthClientInformationFull.model_validate({
|
||||
"client_id": "client",
|
||||
"redirect_uris": ["http://127.0.0.1:47823/callback"],
|
||||
"token_endpoint_auth_method": "none",
|
||||
})
|
||||
)
|
||||
registry = MCPRegistry()
|
||||
|
||||
tools = await registry.get_tools_async([changed])
|
||||
|
||||
assert tools == {}
|
||||
assert registry.needs_auth == {"linear"}
|
||||
assert await storage.get_tokens() is None
|
||||
assert await storage.get_client_info() is None
|
||||
assert await Fingerprint.load("linear") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_login_clears_needs_auth_and_rediscovers(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
registry = MCPRegistry()
|
||||
await registry.get_tools_async([srv])
|
||||
|
||||
async def perform_login(_srv: MCPStreamableHttp, *, on_url: object) -> None:
|
||||
await _save_valid_oauth_state(_srv)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vibe.core.tools.mcp.registry.perform_oauth_login",
|
||||
new=AsyncMock(side_effect=perform_login),
|
||||
),
|
||||
patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http",
|
||||
new=AsyncMock(return_value=[RemoteTool(name="search")]),
|
||||
),
|
||||
):
|
||||
await registry.login("linear", on_url=AsyncMock())
|
||||
|
||||
assert registry.needs_auth == set()
|
||||
assert registry.status()["linear"] == AuthStatus.OK
|
||||
tools = await registry.get_tools_async([srv])
|
||||
assert "linear_search" in tools
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_login_keeps_needs_auth_when_discovery_fails(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
registry = MCPRegistry()
|
||||
await registry.get_tools_async([srv])
|
||||
|
||||
async def perform_login(_srv: MCPStreamableHttp, *, on_url: object) -> None:
|
||||
await _save_valid_oauth_state(_srv)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vibe.core.tools.mcp.registry.perform_oauth_login",
|
||||
new=AsyncMock(side_effect=perform_login),
|
||||
),
|
||||
patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http",
|
||||
new=AsyncMock(side_effect=RuntimeError("boom")),
|
||||
),
|
||||
):
|
||||
with pytest.raises(MCPOAuthLoginFailed):
|
||||
await registry.login("linear", on_url=AsyncMock())
|
||||
|
||||
assert registry.needs_auth == {"linear"}
|
||||
assert registry.status()["linear"] == AuthStatus.NEEDS_AUTH
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_oauth_failure_drops_cached_tools(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
await _save_valid_oauth_state(srv)
|
||||
registry = MCPRegistry()
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http",
|
||||
new=AsyncMock(return_value=[RemoteTool(name="search")]),
|
||||
):
|
||||
tools = await registry.get_tools_async([srv])
|
||||
|
||||
assert "linear_search" in tools
|
||||
assert registry.count_loaded([srv]) == 1
|
||||
|
||||
await registry.mark_oauth_failure("linear")
|
||||
|
||||
assert registry.needs_auth == {"linear"}
|
||||
assert registry.count_loaded([srv]) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logout_deletes_tokens_client_info_and_fingerprint(
|
||||
self, memory_keyring: MemoryKeyring
|
||||
) -> None:
|
||||
srv = _oauth_server()
|
||||
await _save_valid_oauth_state(srv)
|
||||
storage = KeyringTokenStorage(alias="linear")
|
||||
await storage.set_client_info(
|
||||
OAuthClientInformationFull.model_validate({
|
||||
"client_id": "client",
|
||||
"redirect_uris": ["http://127.0.0.1:47823/callback"],
|
||||
"token_endpoint_auth_method": "none",
|
||||
})
|
||||
)
|
||||
registry = MCPRegistry()
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http",
|
||||
new=AsyncMock(return_value=[RemoteTool(name="search")]),
|
||||
):
|
||||
await registry.get_tools_async([srv])
|
||||
|
||||
await registry.logout("linear")
|
||||
|
||||
assert registry.needs_auth == {"linear"}
|
||||
assert await storage.get_tokens() is None
|
||||
assert await Fingerprint.load("linear") is None
|
||||
assert await storage.get_client_info() is None
|
||||
|
||||
|
||||
class TestMCPHttpOAuthProxy:
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth_flow_error_marks_auth_failure_with_stop_turn_message(
|
||||
self,
|
||||
) -> None:
|
||||
callback = AsyncMock()
|
||||
tool_cls = create_mcp_http_proxy_tool_class(
|
||||
url="https://mcp.example.com/mcp",
|
||||
remote=RemoteTool(name="search"),
|
||||
alias="linear",
|
||||
oauth_runtime=MCPHttpOAuthRuntime(
|
||||
lock=asyncio.Lock(), failure_callback=callback
|
||||
),
|
||||
)
|
||||
tool = tool_cls.from_config(lambda: BaseToolConfig())
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.tools.call_tool_http",
|
||||
new=AsyncMock(side_effect=OAuthFlowError("invalid_grant")),
|
||||
):
|
||||
with pytest.raises(ToolError, match="lost authentication"):
|
||||
async for _ in tool.run(_OpenArgs(), InvokeContext(tool_call_id="tc")):
|
||||
pass
|
||||
|
||||
callback.assert_awaited_once_with("linear")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oauth_lock_serializes_concurrent_calls(self) -> None:
|
||||
active = 0
|
||||
max_active = 0
|
||||
|
||||
async def fake_call(*_args: object, **_kwargs: object) -> MCPToolResult:
|
||||
nonlocal active, max_active
|
||||
active += 1
|
||||
max_active = max(max_active, active)
|
||||
await asyncio.sleep(0.01)
|
||||
active -= 1
|
||||
return MCPToolResult(server="linear", tool="search")
|
||||
|
||||
tool_cls = create_mcp_http_proxy_tool_class(
|
||||
url="https://mcp.example.com/mcp",
|
||||
remote=RemoteTool(name="search"),
|
||||
alias="linear",
|
||||
oauth_runtime=MCPHttpOAuthRuntime(
|
||||
lock=asyncio.Lock(), failure_callback=AsyncMock()
|
||||
),
|
||||
)
|
||||
|
||||
async def run_once() -> None:
|
||||
tool = tool_cls.from_config(lambda: BaseToolConfig())
|
||||
async for _ in tool.run(_OpenArgs(), InvokeContext(tool_call_id="tc")):
|
||||
pass
|
||||
|
||||
with patch("vibe.core.tools.mcp.tools.call_tool_http", new=fake_call):
|
||||
await asyncio.gather(run_once(), run_once())
|
||||
|
||||
assert max_active == 1
|
||||
Loading…
Add table
Add a link
Reference in a new issue