f71bfd3b8c
Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Michel Thomazo <51709227+michelTho@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: Mistral Vibe <vibe@mistral.ai>
368 lines
12 KiB
Python
368 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import asyncio.subprocess as aio_subprocess
|
|
import contextlib
|
|
import io
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
from acp import PROTOCOL_VERSION, Client, RequestError, connect_to_agent
|
|
from acp.schema import ClientCapabilities, Implementation
|
|
import pexpect
|
|
import pytest
|
|
|
|
from tests import TESTS_ROOT
|
|
from tests.e2e.common import ansi_tolerant_pattern
|
|
|
|
BROWSER_AUTH_NAME = "Sign in through Mistral AI Studio"
|
|
BROWSER_AUTH_DESCRIPTION = (
|
|
"Sign into Mistral Vibe through your Mistral AI Studio account."
|
|
)
|
|
|
|
|
|
class _AcpSmokeClient(Client):
|
|
def on_connect(self, conn: Any) -> None:
|
|
pass
|
|
|
|
async def request_permission(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "session/request_permission"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def write_text_file(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "fs/write_text_file"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def read_text_file(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "fs/read_text_file"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def create_terminal(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "terminal/create"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def terminal_output(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "terminal/output"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def release_terminal(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "terminal/release"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def wait_for_terminal_exit(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "terminal/wait_for_exit"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def kill_terminal(self, *args: Any, **kwargs: Any) -> Any:
|
|
msg = "terminal/kill"
|
|
raise RequestError.method_not_found(msg)
|
|
|
|
async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
|
|
_ = params
|
|
raise RequestError.method_not_found(method)
|
|
|
|
async def ext_notification(self, method: str, params: dict[str, Any]) -> None:
|
|
_ = params
|
|
raise RequestError.method_not_found(method)
|
|
|
|
async def session_update(self, *_args: Any, **_kwargs: Any) -> None:
|
|
pass
|
|
|
|
|
|
@pytest.fixture
|
|
def vibe_home_dir(tmp_path: Path) -> Path:
|
|
return tmp_path / ".vibe"
|
|
|
|
|
|
async def _spawn_vibe_acp(env: dict[str, str]) -> asyncio.subprocess.Process:
|
|
return await asyncio.create_subprocess_exec(
|
|
"uv",
|
|
"run",
|
|
"vibe-acp",
|
|
stdin=aio_subprocess.PIPE,
|
|
stdout=aio_subprocess.PIPE,
|
|
stderr=aio_subprocess.PIPE,
|
|
cwd=TESTS_ROOT.parent,
|
|
env=env,
|
|
)
|
|
|
|
|
|
async def _terminate_process(proc: asyncio.subprocess.Process) -> None:
|
|
if proc.returncode is None:
|
|
with contextlib.suppress(ProcessLookupError):
|
|
proc.terminate()
|
|
with contextlib.suppress(TimeoutError):
|
|
await asyncio.wait_for(proc.wait(), timeout=5)
|
|
|
|
if proc.returncode is None:
|
|
with contextlib.suppress(ProcessLookupError):
|
|
proc.kill()
|
|
await proc.wait()
|
|
|
|
|
|
def _build_env(
|
|
vibe_home_dir: Path, *, include_api_key: bool, enable_browser_sign_in: bool = False
|
|
) -> dict[str, str]:
|
|
env = os.environ.copy()
|
|
env["PYTHONUNBUFFERED"] = "1"
|
|
env["VIBE_HOME"] = str(vibe_home_dir)
|
|
|
|
vibe_home_dir.mkdir(parents=True, exist_ok=True)
|
|
config_lines = ["enable_telemetry = false"]
|
|
if enable_browser_sign_in:
|
|
config_lines.append("enable_experimental_browser_sign_in = true")
|
|
config_file = vibe_home_dir / "config.toml"
|
|
if not config_file.exists():
|
|
config_file.write_text("\n".join(config_lines) + "\n")
|
|
elif enable_browser_sign_in:
|
|
config_file.write_text(config_file.read_text() + config_lines[-1] + "\n")
|
|
|
|
if include_api_key:
|
|
env["MISTRAL_API_KEY"] = "mock"
|
|
else:
|
|
env.pop("MISTRAL_API_KEY", None)
|
|
|
|
return env
|
|
|
|
|
|
def _build_client_capabilities(
|
|
*, terminal_auth: bool = False, delegated_browser_auth: bool = False
|
|
) -> ClientCapabilities:
|
|
if not terminal_auth and not delegated_browser_auth:
|
|
return ClientCapabilities()
|
|
|
|
field_meta: dict[str, bool] = {}
|
|
if terminal_auth:
|
|
field_meta["terminal-auth"] = True
|
|
if delegated_browser_auth:
|
|
field_meta["browser-auth-delegated"] = True
|
|
return ClientCapabilities(field_meta=field_meta)
|
|
|
|
|
|
async def _connect_and_initialize(
|
|
*,
|
|
vibe_home_dir: Path,
|
|
include_api_key: bool,
|
|
enable_browser_sign_in: bool = False,
|
|
terminal_auth: bool = False,
|
|
delegated_browser_auth: bool = False,
|
|
) -> tuple[asyncio.subprocess.Process, Any, Any]:
|
|
env = _build_env(
|
|
vibe_home_dir,
|
|
include_api_key=include_api_key,
|
|
enable_browser_sign_in=enable_browser_sign_in,
|
|
)
|
|
proc = await _spawn_vibe_acp(env)
|
|
|
|
try:
|
|
assert proc.stdin is not None
|
|
assert proc.stdout is not None
|
|
|
|
conn = connect_to_agent(_AcpSmokeClient(), proc.stdin, proc.stdout)
|
|
initialize_response = await asyncio.wait_for(
|
|
conn.initialize(
|
|
protocol_version=PROTOCOL_VERSION,
|
|
client_capabilities=_build_client_capabilities(
|
|
terminal_auth=terminal_auth,
|
|
delegated_browser_auth=delegated_browser_auth,
|
|
),
|
|
client_info=Implementation(
|
|
name="pytest-smoke", title="Pytest Smoke", version="0.0.0"
|
|
),
|
|
),
|
|
timeout=10,
|
|
)
|
|
except Exception:
|
|
await _terminate_process(proc)
|
|
raise
|
|
|
|
return proc, initialize_response, conn
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_initialize_and_new_session(vibe_home_dir: Path) -> None:
|
|
proc, initialize_response, conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir, include_api_key=True
|
|
)
|
|
|
|
try:
|
|
assert initialize_response.protocol_version == PROTOCOL_VERSION
|
|
assert initialize_response.agent_info.name == "@mistralai/mistral-vibe"
|
|
assert initialize_response.agent_info.title == "Mistral Vibe"
|
|
|
|
session = await asyncio.wait_for(
|
|
conn.new_session(cwd=str(Path.cwd()), mcp_servers=[]), timeout=10
|
|
)
|
|
|
|
assert session.session_id
|
|
finally:
|
|
await _terminate_process(proc)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_bootstraps_default_files(vibe_home_dir: Path) -> None:
|
|
proc, _initialize_response, conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir, include_api_key=True
|
|
)
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
conn.new_session(cwd=str(Path.cwd()), mcp_servers=[]), timeout=10
|
|
)
|
|
finally:
|
|
await _terminate_process(proc)
|
|
assert (vibe_home_dir / "config.toml").is_file()
|
|
assert (vibe_home_dir / "vibehistory").is_file()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_initialize_exposes_browser_auth(vibe_home_dir: Path) -> None:
|
|
proc, initialize_response, _conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir, include_api_key=True, enable_browser_sign_in=True
|
|
)
|
|
|
|
try:
|
|
assert initialize_response.auth_methods is not None
|
|
assert len(initialize_response.auth_methods) == 1
|
|
auth_method = initialize_response.auth_methods[0]
|
|
assert auth_method.id == "browser-auth"
|
|
assert auth_method.name == BROWSER_AUTH_NAME
|
|
assert auth_method.description == BROWSER_AUTH_DESCRIPTION
|
|
finally:
|
|
await _terminate_process(proc)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_initialize_exposes_delegated_browser_auth_when_supported(
|
|
vibe_home_dir: Path,
|
|
) -> None:
|
|
proc, initialize_response, _conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir,
|
|
include_api_key=True,
|
|
enable_browser_sign_in=True,
|
|
delegated_browser_auth=True,
|
|
)
|
|
|
|
try:
|
|
assert initialize_response.auth_methods is not None
|
|
assert len(initialize_response.auth_methods) == 2
|
|
|
|
browser_auth_method = initialize_response.auth_methods[0]
|
|
assert browser_auth_method.id == "browser-auth"
|
|
assert browser_auth_method.name == BROWSER_AUTH_NAME
|
|
assert browser_auth_method.description == BROWSER_AUTH_DESCRIPTION
|
|
|
|
delegated_browser_auth_method = initialize_response.auth_methods[1]
|
|
assert delegated_browser_auth_method.id == "browser-auth-delegated"
|
|
assert delegated_browser_auth_method.name == BROWSER_AUTH_NAME
|
|
assert delegated_browser_auth_method.description == BROWSER_AUTH_DESCRIPTION
|
|
finally:
|
|
await _terminate_process(proc)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_initialize_exposes_terminal_auth_when_supported(
|
|
vibe_home_dir: Path,
|
|
) -> None:
|
|
proc, initialize_response, _conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir,
|
|
include_api_key=True,
|
|
enable_browser_sign_in=True,
|
|
terminal_auth=True,
|
|
)
|
|
|
|
try:
|
|
assert initialize_response.auth_methods is not None
|
|
assert len(initialize_response.auth_methods) == 2
|
|
|
|
browser_auth_method = initialize_response.auth_methods[0]
|
|
assert browser_auth_method.id == "browser-auth"
|
|
assert browser_auth_method.name == BROWSER_AUTH_NAME
|
|
assert browser_auth_method.description == BROWSER_AUTH_DESCRIPTION
|
|
|
|
auth_method = initialize_response.auth_methods[1]
|
|
assert auth_method.id == "vibe-setup"
|
|
assert auth_method.field_meta is not None
|
|
|
|
terminal_auth = auth_method.field_meta["terminal-auth"]
|
|
assert terminal_auth["label"] == "Mistral Vibe Setup"
|
|
assert terminal_auth["command"]
|
|
assert terminal_auth["args"]
|
|
assert terminal_auth["args"][-1:] == ["--setup"]
|
|
finally:
|
|
await _terminate_process(proc)
|
|
|
|
|
|
@pytest.mark.timeout(30)
|
|
def test_vibe_acp_setup_shows_onboarding_and_exits_on_cancel(
|
|
vibe_home_dir: Path,
|
|
) -> None:
|
|
env = cast("os._Environ[str]", _build_env(vibe_home_dir, include_api_key=False))
|
|
env["TERM"] = "xterm-256color"
|
|
|
|
captured = io.StringIO()
|
|
child = pexpect.spawn(
|
|
"uv",
|
|
["run", "vibe-acp", "--setup"],
|
|
cwd=str(TESTS_ROOT.parent),
|
|
env=env,
|
|
encoding="utf-8",
|
|
timeout=10,
|
|
dimensions=(36, 120),
|
|
)
|
|
child.logfile_read = captured
|
|
|
|
try:
|
|
child.expect(ansi_tolerant_pattern("Welcome to"), timeout=15)
|
|
child.sendcontrol("c")
|
|
child.expect(pexpect.EOF, timeout=10)
|
|
finally:
|
|
if child.isalive():
|
|
child.terminate(force=True)
|
|
if not child.closed:
|
|
child.close()
|
|
|
|
output = captured.getvalue()
|
|
assert "Setup cancelled" in output
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_survives_broken_config(vibe_home_dir: Path) -> None:
|
|
vibe_home_dir.mkdir(parents=True, exist_ok=True)
|
|
(vibe_home_dir / "config.toml").write_text("{{{{invalid toml content!!")
|
|
|
|
proc, _initialize_response, conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir, include_api_key=True
|
|
)
|
|
|
|
try:
|
|
# new_session should return a structured JSON-RPC error, not crash the server
|
|
with pytest.raises(RequestError):
|
|
await asyncio.wait_for(
|
|
conn.new_session(cwd=str(Path.cwd()), mcp_servers=[]), timeout=10
|
|
)
|
|
assert proc.returncode is None, "Server crashed after broken config"
|
|
|
|
(vibe_home_dir / "config.toml").write_text("")
|
|
session = await asyncio.wait_for(
|
|
conn.new_session(cwd=str(Path.cwd()), mcp_servers=[]), timeout=10
|
|
)
|
|
assert session.session_id
|
|
finally:
|
|
await _terminate_process(proc)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vibe_acp_new_session_fails_without_api_key(vibe_home_dir: Path) -> None:
|
|
proc, _initialize_response, conn = await _connect_and_initialize(
|
|
vibe_home_dir=vibe_home_dir, include_api_key=False
|
|
)
|
|
|
|
try:
|
|
with pytest.raises(RequestError, match="Missing API key"):
|
|
await asyncio.wait_for(
|
|
conn.new_session(cwd=str(Path.cwd()), mcp_servers=[]), timeout=10
|
|
)
|
|
finally:
|
|
await _terminate_process(proc)
|