mirror of
https://github.com/mistralai/mistral-vibe
synced 2026-04-25 17:14:55 +02:00
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Gauthier Guinet <43207538+Gguinet@users.noreply.github.com> Co-authored-by: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai> Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Quentin <torroba.q@gmail.com> Co-authored-by: Simon <80467011+sorgfresser@users.noreply.github.com> Co-authored-by: Simon Van de Kerckhove <simon.vandekerckhove@mistral.ai> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: angelapopopo <angele.lenglemetz@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
211 lines
8.3 KiB
Python
211 lines
8.3 KiB
Python
from __future__ import annotations
|
||
|
||
import pytest
|
||
|
||
from tests.mock.utils import collect_result
|
||
from vibe.core.tools.base import BaseToolState, ToolError, ToolPermission
|
||
from vibe.core.tools.builtins.bash import Bash, BashArgs, BashToolConfig
|
||
from vibe.core.tools.permissions import PermissionContext
|
||
|
||
|
||
@pytest.fixture
|
||
def bash(tmp_path, monkeypatch):
|
||
monkeypatch.chdir(tmp_path)
|
||
config = BashToolConfig()
|
||
return Bash(config=config, state=BaseToolState())
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_runs_echo_successfully(bash):
|
||
result = await collect_result(bash.run(BashArgs(command="echo hello")))
|
||
|
||
assert result.returncode == 0
|
||
assert result.stdout == "hello\n"
|
||
assert result.stderr == ""
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_fails_cat_command_with_missing_file(bash):
|
||
with pytest.raises(ToolError) as err:
|
||
await collect_result(bash.run(BashArgs(command="cat missing_file.txt")))
|
||
|
||
message = str(err.value)
|
||
assert "Command failed" in message
|
||
assert "Return code: 1" in message
|
||
assert "No such file or directory" in message
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_uses_effective_workdir(tmp_path, monkeypatch):
|
||
monkeypatch.chdir(tmp_path)
|
||
config = BashToolConfig()
|
||
bash_tool = Bash(config=config, state=BaseToolState())
|
||
|
||
result = await collect_result(bash_tool.run(BashArgs(command="pwd")))
|
||
|
||
assert result.stdout.strip() == str(tmp_path)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_handles_timeout(bash):
|
||
with pytest.raises(ToolError) as err:
|
||
await collect_result(bash.run(BashArgs(command="sleep 2", timeout=1)))
|
||
|
||
assert "Command timed out after 1s" in str(err.value)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_truncates_output_to_max_bytes(bash):
|
||
config = BashToolConfig(max_output_bytes=5)
|
||
bash_tool = Bash(config=config, state=BaseToolState())
|
||
|
||
result = await collect_result(
|
||
bash_tool.run(BashArgs(command="printf 'abcdefghij'"))
|
||
)
|
||
|
||
assert result.stdout == "abcde"
|
||
assert result.stderr == ""
|
||
assert result.returncode == 0
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_decodes_non_utf8_bytes(bash):
|
||
result = await collect_result(bash.run(BashArgs(command="printf '\\xff\\xfe'")))
|
||
|
||
# accept both possible encodings, as some shells emit escaped bytes as literal strings
|
||
assert result.stdout in {"<EFBFBD><EFBFBD>", "\xff\xfe", r"\xff\xfe"}
|
||
assert result.stderr == ""
|
||
|
||
|
||
def test_find_not_in_default_allowlist():
|
||
bash_tool = Bash(config=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
|
||
)
|
||
|
||
|
||
def test_resolve_permission():
|
||
config = BashToolConfig(allowlist=["echo", "pwd"], denylist=["rm"])
|
||
bash_tool = Bash(config=config, state=BaseToolState())
|
||
|
||
allowlisted = bash_tool.resolve_permission(BashArgs(command="echo hi"))
|
||
denylisted = bash_tool.resolve_permission(BashArgs(command="rm -rf /tmp"))
|
||
mixed = bash_tool.resolve_permission(BashArgs(command="pwd && whoami"))
|
||
empty = bash_tool.resolve_permission(BashArgs(command=""))
|
||
|
||
assert isinstance(allowlisted, PermissionContext)
|
||
assert allowlisted.permission is ToolPermission.ALWAYS
|
||
assert isinstance(denylisted, PermissionContext)
|
||
assert denylisted.permission is ToolPermission.NEVER
|
||
assert isinstance(mixed, PermissionContext)
|
||
assert mixed.permission is ToolPermission.ASK
|
||
assert any(rp.label == "whoami *" for rp in mixed.required_permissions)
|
||
assert empty is None
|
||
|
||
|
||
class TestResolvePermissionWindowsSyntax:
|
||
"""Verify allowlist/denylist works with Windows-style commands."""
|
||
|
||
def _make_bash(self, **kwargs) -> Bash:
|
||
config = BashToolConfig(**kwargs)
|
||
return Bash(config=config, state=BaseToolState())
|
||
|
||
def test_dir_with_windows_flags_allowlisted(self):
|
||
bash_tool = self._make_bash(allowlist=["dir"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="dir /s /b"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|
||
|
||
def test_type_command_allowlisted(self):
|
||
bash_tool = self._make_bash(allowlist=["type"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="type file.txt"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|
||
|
||
def test_findstr_allowlisted(self):
|
||
bash_tool = self._make_bash(allowlist=["findstr"])
|
||
result = bash_tool.resolve_permission(
|
||
BashArgs(command="findstr /s pattern *.txt")
|
||
)
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|
||
|
||
def test_ver_allowlisted(self):
|
||
bash_tool = self._make_bash(allowlist=["ver"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="ver"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|
||
|
||
def test_where_allowlisted(self):
|
||
bash_tool = self._make_bash(allowlist=["where"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="where python"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|
||
|
||
def test_cmd_k_denylisted(self):
|
||
bash_tool = self._make_bash(denylist=["cmd /k"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="cmd /k something"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.NEVER
|
||
|
||
def test_powershell_noexit_denylisted(self):
|
||
bash_tool = self._make_bash(denylist=["powershell -NoExit"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="powershell -NoExit"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.NEVER
|
||
|
||
def test_notepad_denylisted(self):
|
||
bash_tool = self._make_bash(denylist=["notepad"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="notepad file.txt"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.NEVER
|
||
|
||
def test_cmd_standalone_denylisted(self):
|
||
bash_tool = self._make_bash(denylist_standalone=["cmd"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="cmd"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.NEVER
|
||
|
||
def test_powershell_standalone_denylisted(self):
|
||
bash_tool = self._make_bash(denylist_standalone=["powershell"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="powershell"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.NEVER
|
||
|
||
def test_powershell_cmdlet_asks(self):
|
||
bash_tool = self._make_bash(allowlist=["dir", "echo"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="Get-ChildItem -Path ."))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission == ToolPermission.ASK
|
||
|
||
def test_mixed_allowed_and_unknown_asks(self):
|
||
bash_tool = self._make_bash(allowlist=["git status"])
|
||
result = bash_tool.resolve_permission(
|
||
BashArgs(command="git status && npm install")
|
||
)
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission == ToolPermission.ASK
|
||
|
||
def test_chained_windows_commands_all_allowed(self):
|
||
bash_tool = self._make_bash(allowlist=["dir", "echo"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="dir /s && echo done"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|
||
|
||
def test_chained_commands_one_denied(self):
|
||
bash_tool = self._make_bash(allowlist=["dir"], denylist=["rm"])
|
||
result = bash_tool.resolve_permission(BashArgs(command="dir /s && rm -rf /"))
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.NEVER
|
||
|
||
def test_piped_windows_commands(self):
|
||
bash_tool = self._make_bash(allowlist=["findstr", "type"])
|
||
result = bash_tool.resolve_permission(
|
||
BashArgs(command="type file.txt | findstr pattern")
|
||
)
|
||
assert isinstance(result, PermissionContext)
|
||
assert result.permission is ToolPermission.ALWAYS
|