v2.6.0 (#524)
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>
This commit is contained in:
Mathias Gesbert
GitHub
parent
5103019b01
commit
eb580209d4
180 changed files with 11136 additions and 1030 deletions
|
|
@ -759,6 +759,42 @@ class TestToolCallStructure:
|
|||
)
|
||||
assert rejected_tool_call is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_permission_options_include_granular_labels_for_bash(
|
||||
self, vibe_home_dir: Path
|
||||
) -> None:
|
||||
"""Bash 'npm install foo' should produce granular labels in permission options."""
|
||||
custom_results = [
|
||||
mock_llm_chunk(
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=FunctionCall(
|
||||
name="bash", arguments='{"command":"npm install foo"}'
|
||||
),
|
||||
type="function",
|
||||
index=0,
|
||||
)
|
||||
]
|
||||
),
|
||||
mock_llm_chunk(content="Done"),
|
||||
]
|
||||
mock_env = get_mocking_env(custom_results)
|
||||
async for process in get_acp_agent_loop_process(
|
||||
mock_env=mock_env, vibe_home=vibe_home_dir
|
||||
):
|
||||
permission_request = await start_session_with_request_permission(
|
||||
process, "Run npm install foo"
|
||||
)
|
||||
assert permission_request.params is not None
|
||||
|
||||
# Verify "Allow always" option includes the pattern label
|
||||
allow_always = next(
|
||||
o
|
||||
for o in permission_request.params.options
|
||||
if o.option_id == ToolOption.ALLOW_ALWAYS
|
||||
)
|
||||
assert "npm install *" in allow_always.name
|
||||
|
||||
@pytest.mark.skip(reason="Long running tool call updates are not implemented yet")
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_call_in_progress_update_structure(
|
||||
|
|
|
|||
254
tests/acp/test_acp_entrypoint_smoke.py
Normal file
254
tests/acp/test_acp_entrypoint_smoke.py
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
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
|
||||
|
||||
|
||||
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) -> dict[str, str]:
|
||||
env = os.environ.copy()
|
||||
env["PYTHONUNBUFFERED"] = "1"
|
||||
env["VIBE_HOME"] = str(vibe_home_dir)
|
||||
|
||||
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) -> ClientCapabilities:
|
||||
if not terminal_auth:
|
||||
return ClientCapabilities()
|
||||
|
||||
return ClientCapabilities(field_meta={"terminal-auth": True})
|
||||
|
||||
|
||||
async def _connect_and_initialize(
|
||||
*, vibe_home_dir: Path, include_api_key: bool, terminal_auth: bool = False
|
||||
) -> tuple[asyncio.subprocess.Process, Any, Any]:
|
||||
env = _build_env(vibe_home_dir, include_api_key=include_api_key)
|
||||
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
|
||||
),
|
||||
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_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, terminal_auth=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 == "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"]
|
||||
finally:
|
||||
await _terminate_process(proc)
|
||||
|
||||
|
||||
@pytest.mark.timeout(15)
|
||||
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 Mistral Vibe"), timeout=10)
|
||||
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_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)
|
||||
|
|
@ -28,7 +28,7 @@ class TestACPInitialize:
|
|||
session_capabilities=SessionCapabilities(list=SessionListCapabilities()),
|
||||
)
|
||||
assert response.agent_info == Implementation(
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.5.0"
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.6.0"
|
||||
)
|
||||
|
||||
assert response.auth_methods == []
|
||||
|
|
@ -52,7 +52,7 @@ class TestACPInitialize:
|
|||
session_capabilities=SessionCapabilities(list=SessionListCapabilities()),
|
||||
)
|
||||
assert response.agent_info == Implementation(
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.5.0"
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.6.0"
|
||||
)
|
||||
|
||||
assert response.auth_methods is not None
|
||||
|
|
|
|||
|
|
@ -100,13 +100,11 @@ class TestLoadSession:
|
|||
|
||||
assert response.config_options is not None
|
||||
assert len(response.config_options) == 2
|
||||
assert response.config_options[0].root.id == "mode"
|
||||
assert response.config_options[0].root.category == "mode"
|
||||
assert response.config_options[0].root.current_value == BuiltinAgentName.DEFAULT
|
||||
assert len(response.config_options[0].root.options) == 5
|
||||
mode_option_values = {
|
||||
opt.value for opt in response.config_options[0].root.options
|
||||
}
|
||||
assert response.config_options[0].id == "mode"
|
||||
assert response.config_options[0].category == "mode"
|
||||
assert response.config_options[0].current_value == BuiltinAgentName.DEFAULT
|
||||
assert len(response.config_options[0].options) == 5
|
||||
mode_option_values = {opt.value for opt in response.config_options[0].options}
|
||||
assert mode_option_values == {
|
||||
BuiltinAgentName.DEFAULT,
|
||||
BuiltinAgentName.CHAT,
|
||||
|
|
@ -114,13 +112,11 @@ class TestLoadSession:
|
|||
BuiltinAgentName.PLAN,
|
||||
BuiltinAgentName.ACCEPT_EDITS,
|
||||
}
|
||||
assert response.config_options[1].root.id == "model"
|
||||
assert response.config_options[1].root.category == "model"
|
||||
assert response.config_options[1].root.current_value == "devstral-latest"
|
||||
assert len(response.config_options[1].root.options) == 2
|
||||
model_option_values = {
|
||||
opt.value for opt in response.config_options[1].root.options
|
||||
}
|
||||
assert response.config_options[1].id == "model"
|
||||
assert response.config_options[1].category == "model"
|
||||
assert response.config_options[1].current_value == "devstral-latest"
|
||||
assert len(response.config_options[1].options) == 2
|
||||
model_option_values = {opt.value for opt in response.config_options[1].options}
|
||||
assert model_option_values == {"devstral-latest", "devstral-small"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -103,11 +103,11 @@ class TestACPNewSession:
|
|||
|
||||
# Mode config option
|
||||
mode_config = session_response.config_options[0]
|
||||
assert mode_config.root.id == "mode"
|
||||
assert mode_config.root.category == "mode"
|
||||
assert mode_config.root.current_value == BuiltinAgentName.DEFAULT
|
||||
assert len(mode_config.root.options) == 5
|
||||
mode_option_values = {opt.value for opt in mode_config.root.options}
|
||||
assert mode_config.id == "mode"
|
||||
assert mode_config.category == "mode"
|
||||
assert mode_config.current_value == BuiltinAgentName.DEFAULT
|
||||
assert len(mode_config.options) == 5
|
||||
mode_option_values = {opt.value for opt in mode_config.options}
|
||||
assert mode_option_values == {
|
||||
BuiltinAgentName.DEFAULT,
|
||||
BuiltinAgentName.CHAT,
|
||||
|
|
@ -118,11 +118,11 @@ class TestACPNewSession:
|
|||
|
||||
# Model config option
|
||||
model_config = session_response.config_options[1]
|
||||
assert model_config.root.id == "model"
|
||||
assert model_config.root.category == "model"
|
||||
assert model_config.root.current_value == "devstral-latest"
|
||||
assert len(model_config.root.options) == 2
|
||||
model_option_values = {opt.value for opt in model_config.root.options}
|
||||
assert model_config.id == "model"
|
||||
assert model_config.category == "model"
|
||||
assert model_config.current_value == "devstral-latest"
|
||||
assert len(model_config.options) == 2
|
||||
model_option_values = {opt.value for opt in model_config.options}
|
||||
assert model_option_values == {"devstral-latest", "devstral-small"}
|
||||
|
||||
@pytest.mark.skip(reason="TODO: Fix this test")
|
||||
|
|
|
|||
|
|
@ -83,8 +83,8 @@ class TestACPSetConfigOptionMode:
|
|||
|
||||
# Verify config_options reflect the new state
|
||||
mode_config = response.config_options[0]
|
||||
assert mode_config.root.id == "mode"
|
||||
assert mode_config.root.current_value == BuiltinAgentName.AUTO_APPROVE
|
||||
assert mode_config.id == "mode"
|
||||
assert mode_config.current_value == BuiltinAgentName.AUTO_APPROVE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_config_option_mode_to_plan(
|
||||
|
|
@ -133,8 +133,8 @@ class TestACPSetConfigOptionMode:
|
|||
) # Chat mode auto-approves read-only tools
|
||||
|
||||
mode_config = response.config_options[0]
|
||||
assert mode_config.root.id == "mode"
|
||||
assert mode_config.root.current_value == BuiltinAgentName.CHAT
|
||||
assert mode_config.id == "mode"
|
||||
assert mode_config.current_value == BuiltinAgentName.CHAT
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_config_option_mode_invalid_returns_none(
|
||||
|
|
@ -205,8 +205,8 @@ class TestACPSetConfigOptionModel:
|
|||
|
||||
# Verify config_options reflect the new state
|
||||
model_config = response.config_options[1]
|
||||
assert model_config.root.id == "model"
|
||||
assert model_config.root.current_value == "devstral-small"
|
||||
assert model_config.id == "model"
|
||||
assert model_config.current_value == "devstral-small"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_config_option_model_invalid_returns_none(
|
||||
|
|
|
|||
254
tests/acp/test_usage_update.py
Normal file
254
tests/acp/test_usage_update.py
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from acp.schema import TextContentBlock, UsageUpdate
|
||||
import pytest
|
||||
|
||||
from tests.acp.conftest import _create_acp_agent
|
||||
from tests.conftest import build_test_vibe_config
|
||||
from tests.stubs.fake_backend import FakeBackend
|
||||
from tests.stubs.fake_client import FakeClient
|
||||
from vibe.acp.acp_agent_loop import VibeAcpAgentLoop
|
||||
from vibe.core.agent_loop import AgentLoop
|
||||
from vibe.core.config import SessionLoggingConfig
|
||||
from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role
|
||||
|
||||
|
||||
def _make_backend(prompt_tokens: int = 100, completion_tokens: int = 50) -> FakeBackend:
|
||||
return FakeBackend(
|
||||
LLMChunk(
|
||||
message=LLMMessage(role=Role.assistant, content="Hi"),
|
||||
usage=LLMUsage(
|
||||
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _make_acp_agent(backend: FakeBackend) -> VibeAcpAgentLoop:
|
||||
config = build_test_vibe_config()
|
||||
|
||||
class PatchedAgentLoop(AgentLoop):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **{**kwargs, "backend": backend})
|
||||
self._base_config = config
|
||||
self.agent_manager.invalidate_config()
|
||||
|
||||
patch("vibe.acp.acp_agent_loop.AgentLoop", side_effect=PatchedAgentLoop).start()
|
||||
return _create_acp_agent()
|
||||
|
||||
|
||||
def _get_fake_client(agent: VibeAcpAgentLoop) -> FakeClient:
|
||||
return agent.client # type: ignore[return-value]
|
||||
|
||||
|
||||
def _get_usage_updates(client: FakeClient) -> list[UsageUpdate]:
|
||||
return [
|
||||
update.update
|
||||
for update in client._session_updates
|
||||
if isinstance(update.update, UsageUpdate)
|
||||
]
|
||||
|
||||
|
||||
class TestPromptResponseUsage:
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_returns_usage_in_response(self) -> None:
|
||||
agent = _make_acp_agent(_make_backend(prompt_tokens=100, completion_tokens=50))
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
response = await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert response.usage is not None
|
||||
assert response.usage.input_tokens == 100
|
||||
assert response.usage.output_tokens == 50
|
||||
assert response.usage.total_tokens == 150
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_usage_optional_fields_are_none(self) -> None:
|
||||
agent = _make_acp_agent(_make_backend())
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
response = await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert response.usage is not None
|
||||
assert response.usage.thought_tokens is None
|
||||
assert response.usage.cached_read_tokens is None
|
||||
assert response.usage.cached_write_tokens is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_usage_accumulates_across_turns(self) -> None:
|
||||
backend = _make_backend(prompt_tokens=100, completion_tokens=50)
|
||||
agent = _make_acp_agent(backend)
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
first = await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
second = await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello again")],
|
||||
)
|
||||
|
||||
assert first.usage is not None
|
||||
assert second.usage is not None
|
||||
# Second turn should have strictly more cumulative tokens
|
||||
assert second.usage.input_tokens > first.usage.input_tokens
|
||||
assert second.usage.output_tokens > first.usage.output_tokens
|
||||
assert second.usage.total_tokens > first.usage.total_tokens
|
||||
|
||||
|
||||
class TestUsageUpdateNotification:
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_sends_usage_update(self) -> None:
|
||||
agent = _make_acp_agent(_make_backend())
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
usage_updates = _get_usage_updates(_get_fake_client(agent))
|
||||
assert len(usage_updates) == 1
|
||||
assert usage_updates[0].session_update == "usage_update"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_update_contains_context_window_info(self) -> None:
|
||||
agent = _make_acp_agent(_make_backend(prompt_tokens=100, completion_tokens=50))
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
usage_updates = _get_usage_updates(_get_fake_client(agent))
|
||||
assert len(usage_updates) == 1
|
||||
assert usage_updates[0].size > 0
|
||||
assert usage_updates[0].used > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_update_contains_cost_when_pricing_set(self) -> None:
|
||||
agent = _make_acp_agent(
|
||||
_make_backend(prompt_tokens=1_000_000, completion_tokens=500_000)
|
||||
)
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
# Set pricing directly on the session stats (config loading uses fixture defaults)
|
||||
acp_session = agent.sessions[session.session_id]
|
||||
acp_session.agent_loop.stats.update_pricing(input_price=0.4, output_price=2.0)
|
||||
|
||||
await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
usage_updates = _get_usage_updates(_get_fake_client(agent))
|
||||
assert len(usage_updates) == 1
|
||||
cost = usage_updates[0].cost
|
||||
assert cost is not None
|
||||
assert cost.currency == "USD"
|
||||
assert cost.amount > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_update_no_cost_when_zero_pricing(self) -> None:
|
||||
agent = _make_acp_agent(_make_backend())
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
usage_updates = _get_usage_updates(_get_fake_client(agent))
|
||||
assert len(usage_updates) == 1
|
||||
assert usage_updates[0].cost is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_update_sent_per_prompt(self) -> None:
|
||||
backend = _make_backend()
|
||||
agent = _make_acp_agent(backend)
|
||||
session = await agent.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
||||
|
||||
await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello")],
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
await agent.prompt(
|
||||
session_id=session.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="Hello again")],
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
usage_updates = _get_usage_updates(_get_fake_client(agent))
|
||||
assert len(usage_updates) == 2
|
||||
|
||||
|
||||
class TestLoadSessionUsageUpdate:
|
||||
def _make_session_dir(self, tmp_path: Path, session_id: str, cwd: str) -> Path:
|
||||
session_folder = tmp_path / f"session_20240101_120000_{session_id[:8]}"
|
||||
session_folder.mkdir()
|
||||
messages_file = session_folder / "messages.jsonl"
|
||||
with messages_file.open("w") as f:
|
||||
f.write(json.dumps({"role": "user", "content": "Hello"}) + "\n")
|
||||
meta = {
|
||||
"session_id": session_id,
|
||||
"start_time": "2024-01-01T12:00:00Z",
|
||||
"end_time": "2024-01-01T12:05:00Z",
|
||||
"environment": {"working_directory": cwd},
|
||||
}
|
||||
with (session_folder / "meta.json").open("w") as f:
|
||||
json.dump(meta, f)
|
||||
return session_folder
|
||||
|
||||
def _make_agent_with_session_logging(
|
||||
self, backend: FakeBackend, session_dir: Path
|
||||
) -> VibeAcpAgentLoop:
|
||||
session_config = SessionLoggingConfig(
|
||||
save_dir=str(session_dir), session_prefix="session", enabled=True
|
||||
)
|
||||
config = build_test_vibe_config(session_logging=session_config)
|
||||
|
||||
class PatchedAgentLoop(AgentLoop):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **{**kwargs, "backend": backend})
|
||||
self._base_config = config
|
||||
self.agent_manager.invalidate_config()
|
||||
|
||||
patch("vibe.acp.acp_agent_loop.AgentLoop", side_effect=PatchedAgentLoop).start()
|
||||
agent = _create_acp_agent()
|
||||
patch.object(agent, "_load_config", return_value=config).start()
|
||||
return agent
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_session_sends_usage_update(self, tmp_path: Path) -> None:
|
||||
backend = _make_backend()
|
||||
agent = self._make_agent_with_session_logging(backend, tmp_path)
|
||||
session_id = "test-session-load-usage"
|
||||
self._make_session_dir(tmp_path, session_id, str(Path.cwd()))
|
||||
|
||||
await agent.load_session(cwd=str(Path.cwd()), session_id=session_id)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
client = _get_fake_client(agent)
|
||||
usage_updates = _get_usage_updates(client)
|
||||
assert len(usage_updates) == 1
|
||||
assert usage_updates[0].session_update == "usage_update"
|
||||
assert usage_updates[0].size > 0
|
||||
|
|
@ -1,8 +1,14 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from vibe.acp.utils import get_proxy_help_text
|
||||
from vibe.acp.utils import (
|
||||
TOOL_OPTIONS,
|
||||
ToolOption,
|
||||
build_permission_options,
|
||||
get_proxy_help_text,
|
||||
)
|
||||
from vibe.core.paths import GLOBAL_ENV_FILE
|
||||
from vibe.core.proxy_setup import SUPPORTED_PROXY_VARS
|
||||
from vibe.core.tools.permissions import PermissionScope, RequiredPermission
|
||||
|
||||
|
||||
def _write_env_file(content: str) -> None:
|
||||
|
|
@ -53,3 +59,65 @@ class TestGetProxyHelpText:
|
|||
|
||||
assert "HTTP_PROXY=http://proxy:8080" in result
|
||||
assert "HTTPS_PROXY=" not in result
|
||||
|
||||
|
||||
class TestBuildPermissionOptions:
|
||||
def test_no_permissions_returns_default_options(self) -> None:
|
||||
result = build_permission_options(None)
|
||||
assert result is TOOL_OPTIONS
|
||||
|
||||
def test_empty_list_returns_default_options(self) -> None:
|
||||
result = build_permission_options([])
|
||||
assert result is TOOL_OPTIONS
|
||||
|
||||
def test_with_permissions_includes_labels_in_allow_always(self) -> None:
|
||||
permissions = [
|
||||
RequiredPermission(
|
||||
scope=PermissionScope.COMMAND_PATTERN,
|
||||
invocation_pattern="npm install foo",
|
||||
session_pattern="npm install *",
|
||||
label="npm install *",
|
||||
)
|
||||
]
|
||||
result = build_permission_options(permissions)
|
||||
|
||||
assert len(result) == 3
|
||||
allow_always = next(o for o in result if o.option_id == ToolOption.ALLOW_ALWAYS)
|
||||
assert "npm install *" in allow_always.name
|
||||
assert "session" in allow_always.name.lower()
|
||||
|
||||
def test_allow_always_has_field_meta(self) -> None:
|
||||
permissions = [
|
||||
RequiredPermission(
|
||||
scope=PermissionScope.COMMAND_PATTERN,
|
||||
invocation_pattern="mkdir foo",
|
||||
session_pattern="mkdir *",
|
||||
label="mkdir *",
|
||||
)
|
||||
]
|
||||
result = build_permission_options(permissions)
|
||||
|
||||
allow_always = next(o for o in result if o.option_id == ToolOption.ALLOW_ALWAYS)
|
||||
assert allow_always.field_meta is not None
|
||||
assert "required_permissions" in allow_always.field_meta
|
||||
meta_perms = allow_always.field_meta["required_permissions"]
|
||||
assert len(meta_perms) == 1
|
||||
assert meta_perms[0]["session_pattern"] == "mkdir *"
|
||||
|
||||
def test_allow_once_and_reject_unchanged(self) -> None:
|
||||
permissions = [
|
||||
RequiredPermission(
|
||||
scope=PermissionScope.URL_PATTERN,
|
||||
invocation_pattern="example.com",
|
||||
session_pattern="example.com",
|
||||
label="fetching from example.com",
|
||||
)
|
||||
]
|
||||
result = build_permission_options(permissions)
|
||||
|
||||
allow_once = next(o for o in result if o.option_id == ToolOption.ALLOW_ONCE)
|
||||
reject_once = next(o for o in result if o.option_id == ToolOption.REJECT_ONCE)
|
||||
assert allow_once.name == "Allow once"
|
||||
assert reject_once.name == "Reject once"
|
||||
assert allow_once.field_meta is None
|
||||
assert reject_once.field_meta is None
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue