v2.9.0 (#641)
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:
parent
a83c81ecf5
commit
632ea8c032
253 changed files with 13965 additions and 2525 deletions
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
96
tests/tools/test_scratchpad_permissions.py
Normal file
96
tests/tools/test_scratchpad_permissions.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue