v2.3.0 (#429)
Co-authored-by: Carlo <carloantonio.patti@mistral.ai> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Clement Sirieix <clem.sirieix@gmail.com> Co-authored-by: Michel Thomazo <michel.thomazo@mistral.ai> Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Vincent Guilloux <vincent.guilloux@mistral.ai> Co-authored-by: Thomas Kenbeek <thomas.kenbeek@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Quentin
GitHub
parent
a560a47ce8
commit
5d2e01a6d7
139 changed files with 7152 additions and 1457 deletions
|
|
@ -76,14 +76,14 @@ async def test_decodes_non_utf8_bytes(bash):
|
|||
assert result.stderr == ""
|
||||
|
||||
|
||||
def test_check_allowlist_denylist():
|
||||
def test_resolve_permission():
|
||||
config = BashToolConfig(allowlist=["echo", "pwd"], denylist=["rm"])
|
||||
bash_tool = Bash(config=config, state=BaseToolState())
|
||||
|
||||
allowlisted = bash_tool.check_allowlist_denylist(BashArgs(command="echo hi"))
|
||||
denylisted = bash_tool.check_allowlist_denylist(BashArgs(command="rm -rf /tmp"))
|
||||
mixed = bash_tool.check_allowlist_denylist(BashArgs(command="pwd && whoami"))
|
||||
empty = bash_tool.check_allowlist_denylist(BashArgs(command=""))
|
||||
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 allowlisted is ToolPermission.ALWAYS
|
||||
assert denylisted is ToolPermission.NEVER
|
||||
|
|
|
|||
|
|
@ -5,21 +5,15 @@ import shutil
|
|||
import pytest
|
||||
|
||||
from tests.mock.utils import collect_result
|
||||
from vibe.core.tools.base import ToolError
|
||||
from vibe.core.tools.builtins.grep import (
|
||||
Grep,
|
||||
GrepArgs,
|
||||
GrepBackend,
|
||||
GrepState,
|
||||
GrepToolConfig,
|
||||
)
|
||||
from vibe.core.tools.base import BaseToolState, ToolError
|
||||
from vibe.core.tools.builtins.grep import Grep, GrepArgs, GrepBackend, GrepToolConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def grep(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
config = GrepToolConfig()
|
||||
return Grep(config=config, state=GrepState())
|
||||
return Grep(config=config, state=BaseToolState())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -34,7 +28,7 @@ def grep_gnu_only(tmp_path, monkeypatch):
|
|||
|
||||
monkeypatch.setattr("shutil.which", mock_which)
|
||||
config = GrepToolConfig()
|
||||
return Grep(config=config, state=GrepState())
|
||||
return Grep(config=config, state=BaseToolState())
|
||||
|
||||
|
||||
def test_detects_ripgrep_when_available(grep):
|
||||
|
|
@ -143,7 +137,7 @@ async def test_truncates_to_max_matches(grep, tmp_path):
|
|||
async def test_truncates_to_max_output_bytes(grep, tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
config = GrepToolConfig(max_output_bytes=100)
|
||||
grep_tool = Grep(config=config, state=GrepState())
|
||||
grep_tool = Grep(config=config, state=BaseToolState())
|
||||
(tmp_path / "test.py").write_text("\n".join("x" * 100 for _ in range(10)))
|
||||
|
||||
result = await collect_result(grep_tool.run(GrepArgs(pattern="x")))
|
||||
|
|
@ -191,22 +185,11 @@ async def test_ignores_comments_in_vibeignore(grep, tmp_path):
|
|||
assert result.match_count >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tracks_search_history(grep, tmp_path):
|
||||
(tmp_path / "test.py").write_text("content\n")
|
||||
|
||||
await collect_result(grep.run(GrepArgs(pattern="first")))
|
||||
await collect_result(grep.run(GrepArgs(pattern="second")))
|
||||
await collect_result(grep.run(GrepArgs(pattern="third")))
|
||||
|
||||
assert grep.state.search_history == ["first", "second", "third"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_effective_workdir(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
config = GrepToolConfig()
|
||||
grep_tool = Grep(config=config, state=GrepState())
|
||||
grep_tool = Grep(config=config, state=BaseToolState())
|
||||
(tmp_path / "test.py").write_text("match\n")
|
||||
|
||||
result = await collect_result(grep_tool.run(GrepArgs(pattern="match", path=".")))
|
||||
|
|
|
|||
|
|
@ -227,6 +227,70 @@ class FileTool(BaseTool[FileToolArgs, FileToolResult, BaseToolConfig, BaseToolSt
|
|||
assert "file_tool" in tools
|
||||
|
||||
|
||||
class TestToolRuntimeAvailability:
|
||||
"""Tests for is_available() filtering in ToolManager."""
|
||||
|
||||
def test_unavailable_tool_excluded_from_available_tools(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
):
|
||||
"""Tools where is_available() returns False should be excluded."""
|
||||
import sys
|
||||
|
||||
tool_dir = tmp_path / "tools"
|
||||
tool_dir.mkdir()
|
||||
(tool_dir / "conditional_tool.py").write_text("""
|
||||
import os
|
||||
from vibe.core.tools.base import BaseTool, BaseToolConfig, BaseToolState
|
||||
from pydantic import BaseModel
|
||||
|
||||
class ConditionalToolArgs(BaseModel):
|
||||
pass
|
||||
|
||||
class ConditionalToolResult(BaseModel):
|
||||
pass
|
||||
|
||||
class ConditionalTool(BaseTool[ConditionalToolArgs, ConditionalToolResult, BaseToolConfig, BaseToolState]):
|
||||
description = "Tool that requires TEST_VAR"
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> bool:
|
||||
return bool(os.getenv("TEST_VAR"))
|
||||
|
||||
async def run(self, args, ctx=None):
|
||||
yield ConditionalToolResult()
|
||||
""")
|
||||
|
||||
to_remove = [k for k in sys.modules if "conditional_tool" in k]
|
||||
for k in to_remove:
|
||||
del sys.modules[k]
|
||||
|
||||
monkeypatch.delenv("TEST_VAR", raising=False)
|
||||
vibe_config = build_test_vibe_config(
|
||||
system_prompt_id="tests",
|
||||
include_project_context=False,
|
||||
tool_paths=[tool_dir],
|
||||
)
|
||||
manager = ToolManager(lambda: vibe_config)
|
||||
assert "conditional_tool" not in manager.available_tools
|
||||
|
||||
to_remove = [k for k in sys.modules if "conditional_tool" in k]
|
||||
for k in to_remove:
|
||||
del sys.modules[k]
|
||||
|
||||
monkeypatch.setenv("TEST_VAR", "1")
|
||||
manager2 = ToolManager(lambda: vibe_config)
|
||||
assert "conditional_tool" in manager2.available_tools
|
||||
|
||||
def test_default_is_available_returns_true(self):
|
||||
"""Tools without is_available() override should be available."""
|
||||
vibe_config = build_test_vibe_config(
|
||||
system_prompt_id="tests", include_project_context=False
|
||||
)
|
||||
manager = ToolManager(lambda: vibe_config)
|
||||
|
||||
assert "bash" in manager.available_tools
|
||||
|
||||
|
||||
class TestToolManagerModuleReuse:
|
||||
"""Tests for module reuse across ToolManager instances.
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import pytest
|
|||
|
||||
from vibe.core.config import MCPHttp, MCPStdio, MCPStreamableHttp
|
||||
from vibe.core.tools.mcp import (
|
||||
MCPRegistry,
|
||||
MCPToolResult,
|
||||
RemoteTool,
|
||||
_mcp_stderr_capture,
|
||||
|
|
@ -372,3 +373,181 @@ class TestMCPConfigModels:
|
|||
|
||||
# Trailing special chars become underscores which are then stripped
|
||||
assert config.name == "my_server"
|
||||
|
||||
|
||||
class TestMCPRegistry:
|
||||
def _make_http_server(
|
||||
self, name: str, url: str = "http://localhost:8080"
|
||||
) -> MCPHttp:
|
||||
return MCPHttp(name=name, transport="http", url=url)
|
||||
|
||||
def _make_stdio_server(self, name: str, command: str = "python -m srv") -> MCPStdio:
|
||||
return MCPStdio(name=name, transport="stdio", command=command)
|
||||
|
||||
def test_server_key_is_stable(self):
|
||||
srv = self._make_http_server("s1")
|
||||
registry = MCPRegistry()
|
||||
|
||||
assert registry._server_key(srv) == registry._server_key(srv)
|
||||
|
||||
def test_different_configs_produce_different_keys(self):
|
||||
registry = MCPRegistry()
|
||||
s1 = self._make_http_server("s1", url="http://a:1")
|
||||
s2 = self._make_http_server("s2", url="http://b:2")
|
||||
|
||||
assert registry._server_key(s1) != registry._server_key(s2)
|
||||
|
||||
def test_get_tools_caches_discovery(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_http_server("cached")
|
||||
remote = RemoteTool(name="tool_a", description="A tool")
|
||||
proxy = create_mcp_http_proxy_tool_class(
|
||||
url="http://localhost:8080", remote=remote, alias="cached"
|
||||
)
|
||||
|
||||
key = registry._server_key(srv)
|
||||
registry._cache[key] = {proxy.get_name(): proxy}
|
||||
|
||||
tools = registry.get_tools([srv])
|
||||
assert "cached_tool_a" in tools
|
||||
assert tools["cached_tool_a"] is proxy
|
||||
|
||||
def test_get_tools_returns_empty_for_no_servers(self):
|
||||
registry = MCPRegistry()
|
||||
|
||||
assert registry.get_tools([]) == {}
|
||||
|
||||
def test_clear_drops_cache(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_http_server("s")
|
||||
proxy = create_mcp_http_proxy_tool_class(
|
||||
url="http://localhost:8080", remote=RemoteTool(name="t"), alias="s"
|
||||
)
|
||||
key = registry._server_key(srv)
|
||||
registry._cache[key] = {proxy.get_name(): proxy}
|
||||
|
||||
registry.clear()
|
||||
|
||||
assert len(registry._cache) == 0
|
||||
|
||||
def test_cache_survives_multiple_get_tools_calls(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_http_server("stable")
|
||||
remote = RemoteTool(name="t1")
|
||||
proxy = create_mcp_http_proxy_tool_class(
|
||||
url="http://localhost:8080", remote=remote, alias="stable"
|
||||
)
|
||||
|
||||
key = registry._server_key(srv)
|
||||
registry._cache[key] = {proxy.get_name(): proxy}
|
||||
|
||||
first = registry.get_tools([srv])
|
||||
second = registry.get_tools([srv])
|
||||
|
||||
assert first == second
|
||||
assert first["stable_t1"] is second["stable_t1"]
|
||||
|
||||
def test_disjoint_server_lists_across_agents(self):
|
||||
registry = MCPRegistry()
|
||||
|
||||
srv_x = self._make_http_server("x", url="http://x:1")
|
||||
srv_y = self._make_http_server("y", url="http://y:2")
|
||||
|
||||
proxy_x = create_mcp_http_proxy_tool_class(
|
||||
url="http://x:1", remote=RemoteTool(name="tx"), alias="x"
|
||||
)
|
||||
proxy_y = create_mcp_http_proxy_tool_class(
|
||||
url="http://y:2", remote=RemoteTool(name="ty"), alias="y"
|
||||
)
|
||||
|
||||
registry._cache[registry._server_key(srv_x)] = {proxy_x.get_name(): proxy_x}
|
||||
registry._cache[registry._server_key(srv_y)] = {proxy_y.get_name(): proxy_y}
|
||||
|
||||
agent_a_tools = registry.get_tools([srv_x])
|
||||
agent_b_tools = registry.get_tools([srv_y])
|
||||
|
||||
assert "x_tx" in agent_a_tools
|
||||
assert "y_ty" not in agent_a_tools
|
||||
assert "y_ty" in agent_b_tools
|
||||
assert "x_tx" not in agent_b_tools
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discover_http_success(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_http_server("demo", url="http://demo:9090")
|
||||
remote = RemoteTool(name="hello", description="Hi")
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http", return_value=[remote]
|
||||
):
|
||||
tools = await registry._discover_http(srv)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
name = next(iter(tools))
|
||||
assert name == "demo_hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discover_http_failure_returns_none(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_http_server("fail", url="http://fail:1")
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http",
|
||||
side_effect=ConnectionError("down"),
|
||||
):
|
||||
tools = await registry._discover_http(srv)
|
||||
|
||||
assert tools is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discover_stdio_success(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_stdio_server("local", command="python -m local_srv")
|
||||
remote = RemoteTool(name="run", description="Run it")
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_stdio", return_value=[remote]
|
||||
):
|
||||
tools = await registry._discover_stdio(srv)
|
||||
|
||||
assert tools is not None
|
||||
assert len(tools) == 1
|
||||
name = next(iter(tools))
|
||||
assert name == "local_run"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discover_stdio_failure_returns_none(self):
|
||||
registry = MCPRegistry()
|
||||
srv = self._make_stdio_server("broken")
|
||||
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_stdio",
|
||||
side_effect=OSError("no binary"),
|
||||
):
|
||||
tools = await registry._discover_stdio(srv)
|
||||
|
||||
assert tools is None
|
||||
|
||||
def test_get_tools_discovers_only_uncached(self):
|
||||
registry = MCPRegistry()
|
||||
|
||||
cached_srv = self._make_http_server("cached", url="http://c:1")
|
||||
new_srv = self._make_http_server("new", url="http://n:2")
|
||||
|
||||
cached_proxy = create_mcp_http_proxy_tool_class(
|
||||
url="http://c:1", remote=RemoteTool(name="ct"), alias="cached"
|
||||
)
|
||||
registry._cache[registry._server_key(cached_srv)] = {
|
||||
cached_proxy.get_name(): cached_proxy
|
||||
}
|
||||
|
||||
new_remote = RemoteTool(name="nt")
|
||||
with patch(
|
||||
"vibe.core.tools.mcp.registry.list_tools_http", return_value=[new_remote]
|
||||
):
|
||||
tools = registry.get_tools([cached_srv, new_srv])
|
||||
|
||||
assert "cached_ct" in tools
|
||||
assert "new_nt" in tools
|
||||
assert len(registry._cache) == 2
|
||||
|
|
|
|||
171
tests/tools/test_mcp_sampling.py
Normal file
171
tests/tools/test_mcp_sampling.py
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from mcp.types import (
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
ErrorData,
|
||||
SamplingMessage,
|
||||
TextContent,
|
||||
)
|
||||
import pytest
|
||||
|
||||
from tests.mock.utils import mock_llm_chunk
|
||||
from tests.stubs.fake_backend import FakeBackend
|
||||
from vibe.core.tools.mcp_sampling import (
|
||||
MCPSamplingHandler,
|
||||
_extract_text_content,
|
||||
_map_sampling_messages,
|
||||
)
|
||||
from vibe.core.types import LLMMessage, Role
|
||||
|
||||
|
||||
def _make_config(model_name: str = "test-model") -> MagicMock:
|
||||
config = MagicMock()
|
||||
model = MagicMock()
|
||||
model.name = model_name
|
||||
model.temperature = 0.7
|
||||
config.get_active_model.return_value = model
|
||||
return config
|
||||
|
||||
|
||||
def _make_params(
|
||||
messages: list[SamplingMessage] | None = None,
|
||||
system_prompt: str | None = None,
|
||||
temperature: float | None = None,
|
||||
max_tokens: int = 100,
|
||||
) -> CreateMessageRequestParams:
|
||||
if messages is None:
|
||||
messages = [
|
||||
SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))
|
||||
]
|
||||
return CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
systemPrompt=system_prompt,
|
||||
temperature=temperature,
|
||||
maxTokens=max_tokens,
|
||||
)
|
||||
|
||||
|
||||
class TestExtractTextContent:
|
||||
def test_single_text_block(self) -> None:
|
||||
block = TextContent(type="text", text="hello")
|
||||
assert _extract_text_content(block) == "hello"
|
||||
|
||||
def test_list_of_text_blocks(self) -> None:
|
||||
blocks = [
|
||||
TextContent(type="text", text="a"),
|
||||
TextContent(type="text", text="b"),
|
||||
]
|
||||
assert _extract_text_content(blocks) == "a\nb"
|
||||
|
||||
def test_unsupported_single_block(self) -> None:
|
||||
block = MagicMock(type="image", text=None)
|
||||
assert _extract_text_content(block) == ""
|
||||
|
||||
def test_mixed_blocks_skips_non_text(self) -> None:
|
||||
blocks = [TextContent(type="text", text="keep"), MagicMock(type="image")]
|
||||
assert _extract_text_content(blocks) == "keep"
|
||||
|
||||
|
||||
class TestMapSamplingMessages:
|
||||
def test_maps_user_message(self) -> None:
|
||||
msgs = [
|
||||
SamplingMessage(role="user", content=TextContent(type="text", text="hi"))
|
||||
]
|
||||
result = _map_sampling_messages(msgs)
|
||||
assert len(result) == 1
|
||||
assert result[0].role == Role.user
|
||||
assert result[0].content == "hi"
|
||||
|
||||
def test_maps_assistant_message(self) -> None:
|
||||
msgs = [
|
||||
SamplingMessage(
|
||||
role="assistant", content=TextContent(type="text", text="hello")
|
||||
)
|
||||
]
|
||||
result = _map_sampling_messages(msgs)
|
||||
assert result[0].role == Role.assistant
|
||||
assert result[0].content == "hello"
|
||||
|
||||
def test_maps_multiple_messages(self) -> None:
|
||||
msgs = [
|
||||
SamplingMessage(role="user", content=TextContent(type="text", text="q")),
|
||||
SamplingMessage(
|
||||
role="assistant", content=TextContent(type="text", text="a")
|
||||
),
|
||||
]
|
||||
result = _map_sampling_messages(msgs)
|
||||
assert len(result) == 2
|
||||
assert result[0].role == Role.user
|
||||
assert result[1].role == Role.assistant
|
||||
|
||||
|
||||
class TestMCPSamplingHandler:
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_text_response(self) -> None:
|
||||
chunk = mock_llm_chunk(content="LLM says hi")
|
||||
backend = FakeBackend(chunk)
|
||||
config = _make_config("my-model")
|
||||
handler = MCPSamplingHandler(
|
||||
backend_getter=lambda: backend, config_getter=lambda: config
|
||||
)
|
||||
|
||||
result = await handler(MagicMock(), _make_params())
|
||||
|
||||
assert not isinstance(result, Exception)
|
||||
assert isinstance(result, CreateMessageResult)
|
||||
assert result.role == "assistant"
|
||||
assert result.content.type == "text"
|
||||
assert result.content.text == "LLM says hi"
|
||||
assert result.model == "my-model"
|
||||
assert result.stopReason == "endTurn"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_system_prompt_prepended(self) -> None:
|
||||
chunk = mock_llm_chunk(content="ok")
|
||||
backend = FakeBackend(chunk)
|
||||
config = _make_config()
|
||||
handler = MCPSamplingHandler(
|
||||
backend_getter=lambda: backend, config_getter=lambda: config
|
||||
)
|
||||
|
||||
await handler(MagicMock(), _make_params(system_prompt="Be helpful"))
|
||||
|
||||
sent_messages: list[LLMMessage] = backend.requests_messages[0]
|
||||
assert sent_messages[0].role == Role.system
|
||||
assert sent_messages[0].content == "Be helpful"
|
||||
assert sent_messages[1].role == Role.user
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_calls_backend_with_messages(self) -> None:
|
||||
"""Verify the handler forwards messages to the backend."""
|
||||
chunk = mock_llm_chunk(content="ok")
|
||||
backend = FakeBackend(chunk)
|
||||
config = _make_config()
|
||||
handler = MCPSamplingHandler(
|
||||
backend_getter=lambda: backend, config_getter=lambda: config
|
||||
)
|
||||
|
||||
await handler(MagicMock(), _make_params())
|
||||
|
||||
assert len(backend.requests_messages) == 1
|
||||
sent = backend.requests_messages[0]
|
||||
assert len(sent) == 1
|
||||
assert sent[0].role == Role.user
|
||||
assert sent[0].content == "Hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_error_on_backend_failure(self) -> None:
|
||||
backend = FakeBackend(exception_to_raise=RuntimeError("boom"))
|
||||
config = _make_config()
|
||||
handler = MCPSamplingHandler(
|
||||
backend_getter=lambda: backend, config_getter=lambda: config
|
||||
)
|
||||
|
||||
result = await handler(MagicMock(), _make_params())
|
||||
|
||||
assert isinstance(result, ErrorData)
|
||||
assert result.code == -1
|
||||
assert "boom" in result.message
|
||||
|
|
@ -8,7 +8,7 @@ from tests.conftest import build_test_vibe_config
|
|||
from tests.mock.utils import collect_result
|
||||
from vibe.core.agents.manager import AgentManager
|
||||
from vibe.core.agents.models import BUILTIN_AGENTS, AgentType
|
||||
from vibe.core.tools.base import BaseToolState, InvokeContext, ToolError
|
||||
from vibe.core.tools.base import BaseToolState, InvokeContext, ToolError, ToolPermission
|
||||
from vibe.core.tools.builtins.task import Task, TaskArgs, TaskResult, TaskToolConfig
|
||||
from vibe.core.types import AssistantEvent, LLMMessage, Role
|
||||
|
||||
|
|
@ -76,6 +76,50 @@ class TestTaskToolValidation:
|
|||
assert agent.agent_type == AgentType.SUBAGENT
|
||||
|
||||
|
||||
class TestTaskToolResolvePermission:
|
||||
def test_explore_allowed_by_default(self, task_tool: Task) -> None:
|
||||
args = TaskArgs(task="do something", agent="explore")
|
||||
result = task_tool.resolve_permission(args)
|
||||
assert result == ToolPermission.ALWAYS
|
||||
|
||||
def test_unknown_agent_returns_none(self, task_tool: Task) -> None:
|
||||
args = TaskArgs(task="do something", agent="custom_agent")
|
||||
result = task_tool.resolve_permission(args)
|
||||
assert result is None
|
||||
|
||||
def test_denylist_takes_precedence(self) -> None:
|
||||
config = TaskToolConfig(allowlist=["explore"], denylist=["explore"])
|
||||
tool = Task(config=config, state=BaseToolState())
|
||||
args = TaskArgs(task="do something", agent="explore")
|
||||
result = tool.resolve_permission(args)
|
||||
assert result == ToolPermission.NEVER
|
||||
|
||||
def test_glob_pattern_in_allowlist(self) -> None:
|
||||
config = TaskToolConfig(allowlist=["exp*"])
|
||||
tool = Task(config=config, state=BaseToolState())
|
||||
args = TaskArgs(task="do something", agent="explore")
|
||||
result = tool.resolve_permission(args)
|
||||
assert result == ToolPermission.ALWAYS
|
||||
|
||||
def test_glob_pattern_in_denylist(self) -> None:
|
||||
config = TaskToolConfig(denylist=["danger*"])
|
||||
tool = Task(config=config, state=BaseToolState())
|
||||
args = TaskArgs(task="do something", agent="dangerous_agent")
|
||||
result = tool.resolve_permission(args)
|
||||
assert result == ToolPermission.NEVER
|
||||
|
||||
def test_empty_lists_returns_none(self) -> None:
|
||||
config = TaskToolConfig(allowlist=[], denylist=[])
|
||||
tool = Task(config=config, state=BaseToolState())
|
||||
args = TaskArgs(task="do something", agent="explore")
|
||||
result = tool.resolve_permission(args)
|
||||
assert result is None
|
||||
|
||||
def test_default_config_has_explore_in_allowlist(self) -> None:
|
||||
config = TaskToolConfig()
|
||||
assert "explore" in config.allowlist
|
||||
|
||||
|
||||
class TestTaskToolExecution:
|
||||
@pytest.fixture
|
||||
def ctx(self) -> InvokeContext:
|
||||
|
|
|
|||
253
tests/tools/test_webfetch.py
Normal file
253
tests/tools/test_webfetch.py
Normal file
|
|
@ -0,0 +1,253 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from tests.mock.utils import collect_result
|
||||
from vibe.core.tools.base import BaseToolState, ToolError
|
||||
from vibe.core.tools.builtins.webfetch import WebFetch, WebFetchArgs, WebFetchConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def webfetch():
|
||||
config = WebFetchConfig()
|
||||
return WebFetch(config=config, state=BaseToolState())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def webfetch_small():
|
||||
config = WebFetchConfig(max_content_bytes=100)
|
||||
return WebFetch(config=config, state=BaseToolState())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_bare_domain_gets_https(webfetch):
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text="ok", headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="example.com")))
|
||||
assert result.url == "https://example.com"
|
||||
assert result.content == "ok"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_http_url_stays_http(webfetch):
|
||||
respx.get("http://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text="ok", headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="http://example.com")))
|
||||
assert result.url == "http://example.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_https_url_stays_https(webfetch):
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text="ok", headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
assert result.url == "https://example.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_protocol_relative_url_normalized(webfetch):
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text="ok", headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="//example.com")))
|
||||
assert result.url == "https://example.com"
|
||||
assert result.content == "ok"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ftp_scheme_rejected(webfetch):
|
||||
with pytest.raises(ToolError, match="Invalid URL scheme: ftp"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url="ftp://example.com")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_url_rejected(webfetch):
|
||||
with pytest.raises(ToolError, match="URL cannot be empty"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url=" ")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_html_converted_to_markdown(webfetch):
|
||||
html = "<html><body><h1>Title</h1><p>Hello world</p></body></html>"
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text=html, headers={"Content-Type": "text/html; charset=utf-8"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
assert "# Title" in result.content
|
||||
assert "Hello world" in result.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_plain_text_unchanged(webfetch):
|
||||
respx.get("https://example.com/file.txt").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text="just text", headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(
|
||||
webfetch.run(WebFetchArgs(url="https://example.com/file.txt"))
|
||||
)
|
||||
assert result.content == "just text"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_scripts_stripped_from_markdown(webfetch):
|
||||
html = "<html><body><script>alert('xss')</script><style>.x{}</style><p>Clean</p></body></html>"
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text=html, headers={"Content-Type": "text/html"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
assert "alert" not in result.content
|
||||
assert ".x{}" not in result.content
|
||||
assert "Clean" in result.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_cloudflare_retry_on_challenge(webfetch):
|
||||
route = respx.get("https://example.com")
|
||||
route.side_effect = [
|
||||
httpx.Response(403, headers={"cf-mitigated": "challenge"}),
|
||||
httpx.Response(200, text="success", headers={"Content-Type": "text/plain"}),
|
||||
]
|
||||
result = await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
assert result.content == "success"
|
||||
assert route.call_count == 2
|
||||
|
||||
second_request = route.calls[1].request
|
||||
assert second_request.headers["User-Agent"] == "vibe-cli"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_regular_403_not_retried(webfetch):
|
||||
route = respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(403, headers={"Content-Type": "text/plain"})
|
||||
)
|
||||
with pytest.raises(ToolError, match="HTTP error 403"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_truncates_to_max_bytes_with_disclaimer(webfetch_small):
|
||||
body = "a" * 200
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text=body, headers={"Content-Type": "text/plain"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(
|
||||
webfetch_small.run(WebFetchArgs(url="https://example.com"))
|
||||
)
|
||||
assert result.content.startswith("a" * 100)
|
||||
assert "[Content truncated due to size limit]" in result.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_truncates_html_with_disclaimer(webfetch_small):
|
||||
html = (
|
||||
"<html><body><h2>first title</h2>"
|
||||
+ "x" * 200
|
||||
+ "<h2>second title</h2></body></html>"
|
||||
)
|
||||
respx.get("https://example.com").mock(
|
||||
return_value=httpx.Response(
|
||||
200, text=html, headers={"Content-Type": "text/html"}
|
||||
)
|
||||
)
|
||||
result = await collect_result(
|
||||
webfetch_small.run(WebFetchArgs(url="https://example.com"))
|
||||
)
|
||||
|
||||
assert "## first title" in result.content
|
||||
assert "## second title" not in result.content
|
||||
assert "[Content truncated due to size limit]" in result.content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_http_404_raises_tool_error(webfetch):
|
||||
respx.get("https://example.com").mock(return_value=httpx.Response(404))
|
||||
with pytest.raises(ToolError, match="HTTP error 404"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_http_500_raises_tool_error(webfetch):
|
||||
respx.get("https://example.com").mock(return_value=httpx.Response(500))
|
||||
with pytest.raises(ToolError, match="HTTP error 500"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_timeout_raises_tool_error(webfetch):
|
||||
respx.get("https://example.com").mock(side_effect=httpx.ReadTimeout("timed out"))
|
||||
with pytest.raises(ToolError, match="Request timed out"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_network_error_raises_tool_error(webfetch):
|
||||
respx.get("https://example.com").mock(
|
||||
side_effect=httpx.ConnectError("connection refused")
|
||||
)
|
||||
with pytest.raises(ToolError, match="Failed to fetch URL"):
|
||||
await collect_result(webfetch.run(WebFetchArgs(url="https://example.com")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_negative_timeout_rejected(webfetch):
|
||||
with pytest.raises(ToolError, match="Timeout must be a positive number"):
|
||||
await collect_result(
|
||||
webfetch.run(WebFetchArgs(url="https://example.com", timeout=-1))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_timeout_rejected(webfetch):
|
||||
with pytest.raises(ToolError, match="Timeout must be a positive number"):
|
||||
await collect_result(
|
||||
webfetch.run(WebFetchArgs(url="https://example.com", timeout=0))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_over_max_timeout_rejected(webfetch):
|
||||
with pytest.raises(ToolError, match="Timeout cannot exceed"):
|
||||
await collect_result(
|
||||
webfetch.run(WebFetchArgs(url="https://example.com", timeout=999))
|
||||
)
|
||||
|
||||
|
||||
def test_get_status_text():
|
||||
assert WebFetch.get_status_text() == "Fetching URL"
|
||||
165
tests/tools/test_websearch.py
Normal file
165
tests/tools/test_websearch.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import mistralai
|
||||
import pytest
|
||||
|
||||
from tests.mock.utils import collect_result
|
||||
from vibe.core.tools.base import BaseToolState, ToolError
|
||||
from vibe.core.tools.builtins.websearch import WebSearch, WebSearchArgs, WebSearchConfig
|
||||
|
||||
|
||||
def _make_response(
|
||||
content: list | None = None, outputs: list | None = None
|
||||
) -> mistralai.ConversationResponse:
|
||||
if outputs is None:
|
||||
outputs = [mistralai.MessageOutputEntry(content=content or [])]
|
||||
return mistralai.ConversationResponse(
|
||||
conversation_id="test",
|
||||
outputs=outputs,
|
||||
usage=mistralai.ConversationUsageInfo(
|
||||
prompt_tokens=10, completion_tokens=20, total_tokens=30
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def websearch(monkeypatch):
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "test-key")
|
||||
config = WebSearchConfig()
|
||||
return WebSearch(config=config, state=BaseToolState())
|
||||
|
||||
|
||||
def test_parse_text_chunks(websearch):
|
||||
response = _make_response(
|
||||
content=[mistralai.TextChunk(text="Hello "), mistralai.TextChunk(text="world")]
|
||||
)
|
||||
result = websearch._parse_response(response)
|
||||
assert result.answer == "Hello world"
|
||||
assert result.sources == []
|
||||
|
||||
|
||||
def test_parse_sources_deduped(websearch):
|
||||
response = _make_response(
|
||||
content=[
|
||||
mistralai.TextChunk(text="Answer"),
|
||||
mistralai.ToolReferenceChunk(
|
||||
tool="web_search", title="Site A", url="https://a.com"
|
||||
),
|
||||
mistralai.ToolReferenceChunk(
|
||||
tool="web_search", title="Site A duplicate", url="https://a.com"
|
||||
),
|
||||
mistralai.ToolReferenceChunk(
|
||||
tool="web_search", title="Site B", url="https://b.com"
|
||||
),
|
||||
]
|
||||
)
|
||||
result = websearch._parse_response(response)
|
||||
assert result.answer == "Answer"
|
||||
assert len(result.sources) == 2
|
||||
assert result.sources[0].url == "https://a.com"
|
||||
assert result.sources[0].title == "Site A"
|
||||
assert result.sources[1].url == "https://b.com"
|
||||
|
||||
|
||||
def test_parse_skips_source_without_url(websearch):
|
||||
response = _make_response(
|
||||
content=[
|
||||
mistralai.TextChunk(text="Answer"),
|
||||
mistralai.ToolReferenceChunk(tool="web_search", title="No URL"),
|
||||
]
|
||||
)
|
||||
result = websearch._parse_response(response)
|
||||
assert result.sources == []
|
||||
|
||||
|
||||
def test_parse_empty_text_raises(websearch):
|
||||
response = _make_response(content=[])
|
||||
with pytest.raises(ToolError, match="No text in agent response"):
|
||||
websearch._parse_response(response)
|
||||
|
||||
|
||||
def test_parse_whitespace_only_raises(websearch):
|
||||
response = _make_response(content=[mistralai.TextChunk(text=" ")])
|
||||
with pytest.raises(ToolError, match="No text in agent response"):
|
||||
websearch._parse_response(response)
|
||||
|
||||
|
||||
def test_parse_skips_non_message_entries(websearch):
|
||||
response = _make_response(
|
||||
outputs=[
|
||||
mistralai.MessageOutputEntry(content=[mistralai.TextChunk(text="Answer")])
|
||||
]
|
||||
)
|
||||
result = websearch._parse_response(response)
|
||||
assert result.answer == "Answer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_missing_api_key(monkeypatch):
|
||||
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
|
||||
config = WebSearchConfig()
|
||||
ws = WebSearch(config=config, state=BaseToolState())
|
||||
with pytest.raises(ToolError, match="MISTRAL_API_KEY"):
|
||||
await collect_result(ws.run(WebSearchArgs(query="test")))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_returns_parsed_result(websearch):
|
||||
response = _make_response(
|
||||
content=[
|
||||
mistralai.TextChunk(text="The answer"),
|
||||
mistralai.ToolReferenceChunk(
|
||||
tool="web_search", title="Source", url="https://example.com"
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_start = AsyncMock(return_value=response)
|
||||
with patch.object(mistralai.Mistral, "beta", create=True) as mock_beta:
|
||||
mock_beta.conversations.start_async = mock_start
|
||||
with patch.object(mistralai.Mistral, "__aenter__", return_value=None):
|
||||
with patch.object(mistralai.Mistral, "__aexit__", return_value=None):
|
||||
result = await collect_result(
|
||||
websearch.run(WebSearchArgs(query="test query"))
|
||||
)
|
||||
|
||||
assert result.answer == "The answer"
|
||||
assert len(result.sources) == 1
|
||||
assert result.sources[0].url == "https://example.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_sdk_error_wrapped(websearch):
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
|
||||
mock_response = Mock(spec=httpx.Response)
|
||||
mock_response.status_code = 500
|
||||
mock_response.text = "error"
|
||||
mock_response.headers = httpx.Headers({"content-type": "application/json"})
|
||||
|
||||
with patch.object(mistralai.Mistral, "beta", create=True) as mock_beta:
|
||||
mock_beta.conversations.start_async = AsyncMock(
|
||||
side_effect=mistralai.SDKError("API failed", mock_response)
|
||||
)
|
||||
with patch.object(mistralai.Mistral, "__aenter__", return_value=None):
|
||||
with patch.object(mistralai.Mistral, "__aexit__", return_value=None):
|
||||
with pytest.raises(ToolError, match="Mistral API error"):
|
||||
await collect_result(websearch.run(WebSearchArgs(query="test")))
|
||||
|
||||
|
||||
def test_is_available_with_key(monkeypatch):
|
||||
monkeypatch.setenv("MISTRAL_API_KEY", "key")
|
||||
assert WebSearch.is_available() is True
|
||||
|
||||
|
||||
def test_is_available_without_key(monkeypatch):
|
||||
monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
|
||||
assert WebSearch.is_available() is False
|
||||
|
||||
|
||||
def test_get_status_text():
|
||||
assert WebSearch.get_status_text() == "Searching the web"
|
||||
Loading…
Add table
Add a link
Reference in a new issue