Co-authored-by: Antoine <33425718+anth2o@users.noreply.github.com>
Co-authored-by: Bastien <bastien.baret@gmail.com>
Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai>
Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com>
Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai>
Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai>
Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com>
Co-authored-by: Nelson PROIA <144663685+Nelson-PROIA@users.noreply.github.com>
Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai>
Co-authored-by: Quentin <quentin.torroba@mistral.ai>
Co-authored-by: Robin Gullo <robin.gullo@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin 2026-04-28 17:44:07 +02:00 committed by GitHub
parent a83c81ecf5
commit 632ea8c032
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
253 changed files with 13965 additions and 2525 deletions

View file

@ -77,15 +77,53 @@ async def test_decodes_non_utf8_bytes(bash):
assert result.stderr == ""
def test_find_not_in_default_allowlist():
bash_tool = Bash(config_getter=lambda: BashToolConfig(), state=BaseToolState())
# find -exec runs arbitrary commands; must not be allowlisted by default
permission = bash_tool.resolve_permission(BashArgs(command="find . -exec id \\;"))
assert (
not isinstance(permission, PermissionContext)
or permission.permission is not ToolPermission.ALWAYS
@pytest.mark.parametrize("predicate", ["-exec", "-execdir", "-ok", "-okdir"])
def test_find_execution_predicates_force_ask(predicate: str):
config = BashToolConfig(permission=ToolPermission.ALWAYS)
bash_tool = Bash(config_getter=lambda: config, state=BaseToolState())
permission = bash_tool.resolve_permission(
BashArgs(command=f"find . {predicate} id \\;")
)
assert isinstance(permission, PermissionContext)
assert permission.permission is ToolPermission.ASK
assert [required.label for required in permission.required_permissions] == [
f"find . {predicate} id \\;"
]
def test_find_exec_compound_includes_companion_required_permission():
config = BashToolConfig(permission=ToolPermission.ALWAYS)
bash_tool = Bash(config_getter=lambda: config, state=BaseToolState())
permission = bash_tool.resolve_permission(
BashArgs(command='find . -exec id \\; && python3 -c "import os"')
)
assert isinstance(permission, PermissionContext)
assert permission.permission is ToolPermission.ASK
labels = {rp.label for rp in permission.required_permissions}
assert any("find" in label for label in labels), (
f"Expected a find-exec RequiredPermission, got {labels}"
)
assert any("python3" in label for label in labels), (
f"Companion command should also require permission, got {labels}"
)
def test_find_execution_predicate_does_not_override_denylist():
config = BashToolConfig(denylist=["passwd"])
bash_tool = Bash(config_getter=lambda: config, state=BaseToolState())
permission = bash_tool.resolve_permission(
BashArgs(command="find . -exec id \\; && passwd root")
)
assert isinstance(permission, PermissionContext)
assert permission.permission is ToolPermission.NEVER
assert "matches denylist pattern 'passwd'" in (permission.reason or "")
def test_resolve_permission():
config = BashToolConfig(allowlist=["echo", "pwd"], denylist=["rm"])

View file

@ -9,7 +9,7 @@ import pytest
from tests.stubs.fake_connector_registry import FakeConnectorRegistry
from tests.stubs.fake_mcp_registry import FakeMCPRegistry
from vibe.core.config import VibeConfig
from vibe.core.config import ConnectorConfig, VibeConfig
from vibe.core.tools.base import BaseToolConfig, ToolError
from vibe.core.tools.connectors import CONNECTORS_ENV_VAR
from vibe.core.tools.connectors.connector_registry import (
@ -150,7 +150,7 @@ class TestFakeConnectorRegistry:
class TestToolManagerConnectorIntegration:
@staticmethod
def _make_config() -> VibeConfig:
def _make_config(connectors: list[ConnectorConfig] | None = None) -> VibeConfig:
"""Minimal VibeConfig-like stub for ToolManager."""
return cast(
VibeConfig,
@ -159,6 +159,7 @@ class TestToolManagerConnectorIntegration:
(),
{
"mcp_servers": [],
"connectors": connectors or [],
"enabled_tools": [],
"disabled_tools": [],
"tools": {},
@ -399,3 +400,100 @@ class TestConnectorProxyToolRun:
pytest.raises(ToolError, match="network"),
):
[_ async for _ in tool.invoke(query="hello")]
# ---------------------------------------------------------------------------
# ToolManager: per-connector disabled / disabled_tools filtering
# ---------------------------------------------------------------------------
class TestConnectorDisableFiltering:
@staticmethod
def _make_config(connectors: list[ConnectorConfig] | None = None) -> VibeConfig:
return cast(
VibeConfig,
type(
"_Cfg",
(),
{
"mcp_servers": [],
"connectors": connectors or [],
"enabled_tools": [],
"disabled_tools": [],
"tools": {},
"tool_paths": [],
},
)(),
)
def test_disabled_connector_excludes_all_tools(self) -> None:
registry = FakeConnectorRegistry(
connectors={
"wiki": [
RemoteTool(name="search", description="Search"),
RemoteTool(name="read", description="Read"),
]
}
)
config = self._make_config(
connectors=[ConnectorConfig(name="wiki", disabled=True)]
)
tm = ToolManager(
config_getter=lambda: config,
mcp_registry=FakeMCPRegistry(),
connector_registry=registry,
)
assert "connector_wiki_search" not in tm.available_tools
assert "connector_wiki_read" not in tm.available_tools
# But still registered (discoverable for UI)
assert "connector_wiki_search" in tm.registered_tools
def test_disabled_tools_filters_specific_tools(self) -> None:
registry = FakeConnectorRegistry(
connectors={
"mail": [
RemoteTool(name="send", description="Send"),
RemoteTool(name="read", description="Read"),
]
}
)
config = self._make_config(
connectors=[ConnectorConfig(name="mail", disabled_tools=["send"])]
)
tm = ToolManager(
config_getter=lambda: config,
mcp_registry=FakeMCPRegistry(),
connector_registry=registry,
)
assert "connector_mail_send" not in tm.available_tools
assert "connector_mail_read" in tm.available_tools
def test_no_config_means_all_enabled(self) -> None:
registry = FakeConnectorRegistry(
connectors={"wiki": [RemoteTool(name="search", description="Search")]}
)
config = self._make_config(connectors=[])
tm = ToolManager(
config_getter=lambda: config,
mcp_registry=FakeMCPRegistry(),
connector_registry=registry,
)
assert "connector_wiki_search" in tm.available_tools
def test_unrelated_config_does_not_affect_other_connectors(self) -> None:
registry = FakeConnectorRegistry(
connectors={
"wiki": [RemoteTool(name="search", description="Search")],
"mail": [RemoteTool(name="send", description="Send")],
}
)
config = self._make_config(
connectors=[ConnectorConfig(name="mail", disabled=True)]
)
tm = ToolManager(
config_getter=lambda: config,
mcp_registry=FakeMCPRegistry(),
connector_registry=registry,
)
assert "connector_wiki_search" in tm.available_tools
assert "connector_mail_send" not in tm.available_tools

View file

@ -185,6 +185,12 @@ class TestBashGranularPermissions:
assert isinstance(result, PermissionContext)
assert result.permission is ToolPermission.ASK
def test_sensitive_bypasses_global_always_permission(self):
bash = self._bash(permission=ToolPermission.ALWAYS)
result = bash.resolve_permission(BashArgs(command="sudo ls"))
assert isinstance(result, PermissionContext)
assert result.permission is ToolPermission.ASK
def test_allowlisted_outside_dir_still_asks(self):
bash = self._bash()
# cat is allowlisted but /etc/passwd is outside workdir

View file

@ -10,7 +10,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
from pydantic import ValidationError
import pytest
from vibe.core.config import MCPHttp, MCPStdio, MCPStreamableHttp
from tests.stubs.fake_mcp_registry import FakeMCPRegistry
from vibe.core.config import MCPHttp, MCPStdio, MCPStreamableHttp, VibeConfig
from vibe.core.tools.mcp import (
MCPRegistry,
MCPToolResult,
@ -710,3 +711,140 @@ class TestMCPStdioCwd:
)
assert proxy_cls._cwd is None
# ---------------------------------------------------------------------------
# _MCPBase disabled / disabled_tools field tests
# ---------------------------------------------------------------------------
class TestMCPBaseDisableFields:
def test_disabled_defaults_to_false(self):
config = MCPStdio(name="test", transport="stdio", command="python")
assert config.disabled is False
assert config.disabled_tools == []
def test_disabled_true(self):
config = MCPStdio(
name="test", transport="stdio", command="python", disabled=True
)
assert config.disabled is True
def test_disabled_tools_list(self):
config = MCPHttp(
name="test",
transport="http",
url="http://localhost:8080",
disabled_tools=["search", "read"],
)
assert config.disabled_tools == ["search", "read"]
def test_disabled_fields_on_streamable_http(self):
config = MCPStreamableHttp(
name="test",
transport="streamable-http",
url="http://localhost:8080",
disabled=True,
disabled_tools=["write"],
)
assert config.disabled is True
assert config.disabled_tools == ["write"]
# ---------------------------------------------------------------------------
# ToolManager: per-MCP-server disabled / disabled_tools filtering
# ---------------------------------------------------------------------------
from vibe.core.tools.manager import ToolManager
class TestMCPDisableFiltering:
@staticmethod
def _make_config(
mcp_servers: list[MCPHttp | MCPStdio | MCPStreamableHttp] | None = None,
) -> VibeConfig:
return cast(
VibeConfig,
type(
"_Cfg",
(),
{
"mcp_servers": mcp_servers or [],
"connectors": [],
"enabled_tools": [],
"disabled_tools": [],
"tools": {},
"tool_paths": [],
},
)(),
)
def test_disabled_server_excludes_all_tools(self):
srv = MCPHttp(
name="demo", transport="http", url="http://demo:9090", disabled=True
)
registry = FakeMCPRegistry()
remote_a = RemoteTool(name="tool_a", description="A")
remote_b = RemoteTool(name="tool_b", description="B")
proxy_a = create_mcp_http_proxy_tool_class(
url="http://demo:9090", remote=remote_a, alias="demo"
)
proxy_b = create_mcp_http_proxy_tool_class(
url="http://demo:9090", remote=remote_b, alias="demo"
)
registry.set_tools(
[srv], {proxy_a.get_name(): proxy_a, proxy_b.get_name(): proxy_b}
)
config = self._make_config(mcp_servers=[srv])
tm = ToolManager(
config_getter=lambda: config, mcp_registry=registry, connector_registry=None
)
assert "demo_tool_a" not in tm.available_tools
assert "demo_tool_b" not in tm.available_tools
# Still registered (discoverable for UI)
assert "demo_tool_a" in tm.registered_tools
def test_disabled_tools_filters_specific_tools(self):
srv = MCPHttp(
name="demo",
transport="http",
url="http://demo:9090",
disabled_tools=["tool_a"],
)
registry = FakeMCPRegistry()
remote_a = RemoteTool(name="tool_a", description="A")
remote_b = RemoteTool(name="tool_b", description="B")
proxy_a = create_mcp_http_proxy_tool_class(
url="http://demo:9090", remote=remote_a, alias="demo"
)
proxy_b = create_mcp_http_proxy_tool_class(
url="http://demo:9090", remote=remote_b, alias="demo"
)
registry.set_tools(
[srv], {proxy_a.get_name(): proxy_a, proxy_b.get_name(): proxy_b}
)
config = self._make_config(mcp_servers=[srv])
tm = ToolManager(
config_getter=lambda: config, mcp_registry=registry, connector_registry=None
)
assert "demo_tool_a" not in tm.available_tools
assert "demo_tool_b" in tm.available_tools
def test_disabled_false_is_noop(self):
srv = MCPHttp(
name="demo", transport="http", url="http://demo:9090", disabled=False
)
registry = FakeMCPRegistry()
remote = RemoteTool(name="tool_a", description="A")
proxy = create_mcp_http_proxy_tool_class(
url="http://demo:9090", remote=remote, alias="demo"
)
registry.set_tools([srv], {proxy.get_name(): proxy})
config = self._make_config(mcp_servers=[srv])
tm = ToolManager(
config_getter=lambda: config, mcp_registry=registry, connector_registry=None
)
assert "demo_tool_a" in tm.available_tools

View file

@ -156,6 +156,23 @@ class TestMCPSamplingHandler:
assert sent[0].role == Role.user
assert sent[0].content == "Hello"
@pytest.mark.asyncio
async def test_forwards_metadata_and_headers(self) -> None:
chunk = mock_llm_chunk(content="ok")
backend = FakeBackend(chunk)
config = _make_config()
handler = MCPSamplingHandler(
backend_getter=lambda: backend,
config_getter=lambda: config,
metadata_getter=lambda: {"call_type": "secondary_call"},
extra_headers_getter=lambda: {"x-affinity": "session-123"},
)
await handler(MagicMock(), _make_params())
assert backend.requests_metadata == [{"call_type": "secondary_call"}]
assert backend.requests_extra_headers == [{"x-affinity": "session-123"}]
@pytest.mark.asyncio
async def test_returns_error_on_backend_failure(self) -> None:
backend = FakeBackend(exception_to_raise=RuntimeError("boom"))

View file

@ -0,0 +1,96 @@
from __future__ import annotations
import pytest
import vibe.core.scratchpad as scratchpad_mod
from vibe.core.scratchpad import init_scratchpad
from vibe.core.tools.base import BaseToolState, ToolPermission
from vibe.core.tools.builtins.bash import (
Bash,
BashArgs,
BashToolConfig,
_collect_outside_dirs,
)
from vibe.core.tools.builtins.read_file import (
ReadFile,
ReadFileArgs,
ReadFileState,
ReadFileToolConfig,
)
from vibe.core.tools.builtins.write_file import (
WriteFile,
WriteFileArgs,
WriteFileConfig,
)
from vibe.core.tools.permissions import PermissionContext, PermissionScope
@pytest.fixture(autouse=True)
def _setup_scratchpad():
init_scratchpad("test-session")
class TestFileToolScratchpadPermissions:
def test_write_file_scratchpad_always_allowed(self):
sp = scratchpad_mod.get_scratchpad_dir("test-session")
assert sp is not None
tool = WriteFile(config_getter=lambda: WriteFileConfig(), state=BaseToolState())
result = tool.resolve_permission(
WriteFileArgs(path=str(sp / "draft.py"), content="x")
)
assert isinstance(result, PermissionContext)
assert result.permission is ToolPermission.ALWAYS
def test_read_file_scratchpad_always_allowed(self):
sp = scratchpad_mod.get_scratchpad_dir("test-session")
assert sp is not None
tool = ReadFile(
config_getter=lambda: ReadFileToolConfig(), state=ReadFileState()
)
result = tool.resolve_permission(ReadFileArgs(path=str(sp / "notes.txt")))
assert isinstance(result, PermissionContext)
assert result.permission is ToolPermission.ALWAYS
def test_scratchpad_env_file_still_allowed(self):
"""Scratchpad bypasses sensitive pattern checks."""
sp = scratchpad_mod.get_scratchpad_dir("test-session")
assert sp is not None
tool = WriteFile(config_getter=lambda: WriteFileConfig(), state=BaseToolState())
result = tool.resolve_permission(
WriteFileArgs(path=str(sp / ".env"), content="SECRET=x")
)
assert isinstance(result, PermissionContext)
assert result.permission is ToolPermission.ALWAYS
def test_non_scratchpad_outside_dir_still_asks(self):
tool = WriteFile(config_getter=lambda: WriteFileConfig(), state=BaseToolState())
result = tool.resolve_permission(
WriteFileArgs(path="/tmp/not-scratchpad/file.txt", content="x")
)
assert isinstance(result, PermissionContext)
assert result.permission is ToolPermission.ASK
class TestBashScratchpadPermissions:
def test_scratchpad_path_not_flagged_as_outside_dir(self):
sp = scratchpad_mod.get_scratchpad_dir("test-session")
assert sp is not None
dirs = _collect_outside_dirs([f"cat {sp}/file.txt"])
assert len(dirs) == 0
def test_non_scratchpad_outside_path_still_flagged(self):
dirs = _collect_outside_dirs(["cat /etc/hosts"])
assert len(dirs) >= 1
def test_bash_scratchpad_mkdir_no_outside_dir_permission(self):
sp = scratchpad_mod.get_scratchpad_dir("test-session")
assert sp is not None
bash = Bash(config_getter=lambda: BashToolConfig(), state=BaseToolState())
result = bash.resolve_permission(BashArgs(command=f"mkdir {sp}/subdir"))
assert isinstance(result, PermissionContext)
outside = [
rp
for rp in result.required_permissions
if rp.scope is PermissionScope.OUTSIDE_DIRECTORY
]
assert len(outside) == 0