v2.7.1 (#551)
Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Jean-Malo Delignon <56539593+jean-malo@users.noreply.github.com> Co-authored-by: Paul Cacheux <paul.cacheux@mistral.ai> Co-authored-by: Quentin <torroba.q@gmail.com> Co-authored-by: angelapopopo <angele.lenglemetz@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin
GitHub
parent
6a50d1d521
commit
54b9a17457
30 changed files with 1000 additions and 200 deletions
|
|
@ -463,14 +463,6 @@ class TestSessionUpdates:
|
|||
== "available_commands_update"
|
||||
)
|
||||
|
||||
user_response_text = await read_response(process)
|
||||
assert user_response_text is not None
|
||||
user_response = UpdateJsonRpcNotification.model_validate(
|
||||
json.loads(user_response_text)
|
||||
)
|
||||
assert user_response.params is not None
|
||||
assert user_response.params.update.session_update == "user_message_chunk"
|
||||
|
||||
text_response = await read_response(process)
|
||||
assert text_response is not None
|
||||
response = UpdateJsonRpcNotification.model_validate(
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from __future__ import annotations
|
|||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from acp.schema import AgentThoughtChunk, TextContentBlock
|
||||
from acp.schema import AgentMessageChunk, AgentThoughtChunk, TextContentBlock
|
||||
import pytest
|
||||
|
||||
from tests.acp.conftest import _create_acp_agent
|
||||
|
|
@ -146,9 +146,16 @@ class TestACPAgentThought:
|
|||
for update in fake_client._session_updates
|
||||
if isinstance(update.update, AgentThoughtChunk)
|
||||
]
|
||||
agent_updates = [
|
||||
update
|
||||
for update in fake_client._session_updates
|
||||
if isinstance(update.update, AgentMessageChunk)
|
||||
]
|
||||
|
||||
assert len(thought_updates) == 1
|
||||
thought_chunk = thought_updates[0].update
|
||||
assert thought_chunk.field_meta is not None
|
||||
assert "messageId" in thought_chunk.field_meta
|
||||
assert thought_chunk.field_meta["messageId"] is not None
|
||||
assert thought_chunk.message_id is not None
|
||||
|
||||
assert len(agent_updates) == 1
|
||||
agent_chunk = agent_updates[0].update
|
||||
assert thought_chunk.message_id != agent_chunk.message_id
|
||||
|
|
|
|||
|
|
@ -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.7.0"
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.7.1"
|
||||
)
|
||||
|
||||
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.7.0"
|
||||
name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.7.1"
|
||||
)
|
||||
|
||||
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
|
||||
|
|
@ -342,3 +338,99 @@ class TestLoadSession:
|
|||
assert user_updates[1].update.content.text == "Second message"
|
||||
assert agent_updates[0].update.content.text == "First response"
|
||||
assert agent_updates[1].update.content.text == "Second response"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replay_user_message_has_message_id(
|
||||
self,
|
||||
acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient],
|
||||
temp_session_dir: Path,
|
||||
create_test_session,
|
||||
) -> None:
|
||||
acp_agent, client = acp_agent_with_session_config
|
||||
|
||||
session_id = "msg-id-usr-1234567"
|
||||
cwd = str(Path.cwd())
|
||||
message_id = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
|
||||
messages = [{"role": "user", "content": "Hello", "message_id": message_id}]
|
||||
create_test_session(temp_session_dir, session_id, cwd, messages=messages)
|
||||
|
||||
await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id)
|
||||
|
||||
user_updates = [
|
||||
u for u in client._session_updates if isinstance(u.update, UserMessageChunk)
|
||||
]
|
||||
assert len(user_updates) == 1
|
||||
assert user_updates[0].update.message_id == message_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replay_agent_message_has_message_id(
|
||||
self,
|
||||
acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient],
|
||||
temp_session_dir: Path,
|
||||
create_test_session,
|
||||
) -> None:
|
||||
acp_agent, client = acp_agent_with_session_config
|
||||
|
||||
session_id = "msg-id-ast-1234567"
|
||||
cwd = str(Path.cwd())
|
||||
message_id = "11111111-2222-3333-4444-555555555555"
|
||||
messages = [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello!", "message_id": message_id},
|
||||
]
|
||||
create_test_session(temp_session_dir, session_id, cwd, messages=messages)
|
||||
|
||||
await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id)
|
||||
|
||||
agent_updates = [
|
||||
u
|
||||
for u in client._session_updates
|
||||
if isinstance(u.update, AgentMessageChunk)
|
||||
]
|
||||
assert len(agent_updates) == 1
|
||||
assert agent_updates[0].update.message_id == message_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replay_reasoning_has_different_message_id_than_agent_message(
|
||||
self,
|
||||
acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient],
|
||||
temp_session_dir: Path,
|
||||
create_test_session,
|
||||
) -> None:
|
||||
acp_agent, client = acp_agent_with_session_config
|
||||
|
||||
session_id = "msg-id-rsn-1234567"
|
||||
cwd = str(Path.cwd())
|
||||
agent_message_id = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
|
||||
reasoning_message_id = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
|
||||
messages = [
|
||||
{"role": "user", "content": "Think about this"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Here is my answer",
|
||||
"message_id": agent_message_id,
|
||||
"reasoning_content": "Let me think...",
|
||||
"reasoning_message_id": reasoning_message_id,
|
||||
},
|
||||
]
|
||||
create_test_session(temp_session_dir, session_id, cwd, messages=messages)
|
||||
|
||||
await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id)
|
||||
|
||||
agent_updates = [
|
||||
u
|
||||
for u in client._session_updates
|
||||
if isinstance(u.update, AgentMessageChunk)
|
||||
]
|
||||
thought_updates = [
|
||||
u
|
||||
for u in client._session_updates
|
||||
if isinstance(u.update, AgentThoughtChunk)
|
||||
]
|
||||
assert len(agent_updates) == 1
|
||||
assert len(thought_updates) == 1
|
||||
assert agent_updates[0].update.message_id == agent_message_id
|
||||
assert thought_updates[0].update.message_id == reasoning_message_id
|
||||
assert (
|
||||
agent_updates[0].update.message_id != thought_updates[0].update.message_id
|
||||
)
|
||||
|
|
|
|||
156
tests/acp/test_message_id.py
Normal file
156
tests/acp/test_message_id.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from uuid import UUID
|
||||
|
||||
from acp.schema import AgentMessageChunk, TextContentBlock
|
||||
import pytest
|
||||
|
||||
from tests.acp.conftest import _create_acp_agent
|
||||
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.types import LLMChunk, LLMMessage, LLMUsage, Role
|
||||
|
||||
|
||||
def _is_uuid(value: str) -> bool:
|
||||
try:
|
||||
UUID(value)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _make_response_chunk(content: str = "Hi") -> LLMChunk:
|
||||
return LLMChunk(
|
||||
message=LLMMessage(role=Role.assistant, content=content),
|
||||
usage=LLMUsage(prompt_tokens=1, completion_tokens=1),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def two_turn_acp_agent_loop() -> VibeAcpAgentLoop:
|
||||
backend = FakeBackend([
|
||||
[_make_response_chunk("Hi")],
|
||||
[_make_response_chunk("Hi again")],
|
||||
])
|
||||
|
||||
class PatchedAgentLoop(AgentLoop):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **{**kwargs, "backend": backend})
|
||||
|
||||
patch("vibe.acp.acp_agent_loop.AgentLoop", side_effect=PatchedAgentLoop).start()
|
||||
return _create_acp_agent()
|
||||
|
||||
|
||||
class TestPromptResponseUserMessageId:
|
||||
@pytest.mark.asyncio
|
||||
async def test_generates_user_message_id_when_client_provides_none(
|
||||
self, acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
|
||||
response = await acp_agent_loop.prompt(
|
||||
session_id=session_response.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="hi")],
|
||||
)
|
||||
|
||||
assert response.user_message_id is not None
|
||||
assert _is_uuid(response.user_message_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_echoes_client_provided_message_id(
|
||||
self, acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
client_message_id = "550e8400-e29b-41d4-a716-446655440000"
|
||||
|
||||
response = await acp_agent_loop.prompt(
|
||||
session_id=session_response.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="hi")],
|
||||
message_id=client_message_id,
|
||||
)
|
||||
|
||||
assert response.user_message_id == client_message_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_message_ids_are_unique_across_turns(
|
||||
self, acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
session_id = session_response.session_id
|
||||
|
||||
response_1 = await acp_agent_loop.prompt(
|
||||
session_id=session_id, prompt=[TextContentBlock(type="text", text="hi")]
|
||||
)
|
||||
response_2 = await acp_agent_loop.prompt(
|
||||
session_id=session_id,
|
||||
prompt=[TextContentBlock(type="text", text="hi again")],
|
||||
)
|
||||
|
||||
assert response_1.user_message_id != response_2.user_message_id
|
||||
|
||||
|
||||
class TestAgentMessageChunkMessageId:
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_message_chunk_has_message_id(
|
||||
self, acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
|
||||
await acp_agent_loop.prompt(
|
||||
session_id=session_response.session_id,
|
||||
prompt=[TextContentBlock(type="text", text="hi")],
|
||||
)
|
||||
|
||||
fake_client: FakeClient = acp_agent_loop.client # type: ignore
|
||||
agent_chunks = [
|
||||
u
|
||||
for u in fake_client._session_updates
|
||||
if isinstance(u.update, AgentMessageChunk)
|
||||
]
|
||||
|
||||
assert len(agent_chunks) >= 1
|
||||
assert agent_chunks[0].update.message_id is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_message_ids_are_unique_across_turns(
|
||||
self, two_turn_acp_agent_loop: VibeAcpAgentLoop
|
||||
) -> None:
|
||||
session_response = await two_turn_acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
session_id = session_response.session_id
|
||||
fake_client: FakeClient = two_turn_acp_agent_loop.client # type: ignore
|
||||
|
||||
await two_turn_acp_agent_loop.prompt(
|
||||
session_id=session_id, prompt=[TextContentBlock(type="text", text="hi")]
|
||||
)
|
||||
chunks_turn_1 = [
|
||||
u
|
||||
for u in fake_client._session_updates
|
||||
if isinstance(u.update, AgentMessageChunk)
|
||||
]
|
||||
fake_client._session_updates.clear()
|
||||
|
||||
await two_turn_acp_agent_loop.prompt(
|
||||
session_id=session_id,
|
||||
prompt=[TextContentBlock(type="text", text="hi again")],
|
||||
)
|
||||
chunks_turn_2 = [
|
||||
u
|
||||
for u in fake_client._session_updates
|
||||
if isinstance(u.update, AgentMessageChunk)
|
||||
]
|
||||
|
||||
assert chunks_turn_1[0].update.message_id != chunks_turn_2[0].update.message_id
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -148,6 +148,98 @@ class TestProxySetupCommand:
|
|||
assert "HTTP_PROXY" in env_content
|
||||
assert "http://localhost:8080" in env_content
|
||||
|
||||
|
||||
class TestProxySetupMessageId:
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_setup_response_has_user_message_id(
|
||||
self,
|
||||
acp_agent_loop: VibeAcpAgentLoop,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
env_file = tmp_path / ".env"
|
||||
|
||||
class FakeGlobalEnvFile:
|
||||
path = env_file
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.proxy_setup.GLOBAL_ENV_FILE", FakeGlobalEnvFile()
|
||||
)
|
||||
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
|
||||
response = await acp_agent_loop.prompt(
|
||||
prompt=[TextContentBlock(type="text", text="/proxy-setup")],
|
||||
session_id=session_response.session_id,
|
||||
)
|
||||
|
||||
assert response.user_message_id is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_setup_echoes_client_message_id(
|
||||
self,
|
||||
acp_agent_loop: VibeAcpAgentLoop,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
env_file = tmp_path / ".env"
|
||||
|
||||
class FakeGlobalEnvFile:
|
||||
path = env_file
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.proxy_setup.GLOBAL_ENV_FILE", FakeGlobalEnvFile()
|
||||
)
|
||||
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
client_message_id = "550e8400-e29b-41d4-a716-446655440000"
|
||||
|
||||
response = await acp_agent_loop.prompt(
|
||||
prompt=[TextContentBlock(type="text", text="/proxy-setup")],
|
||||
session_id=session_response.session_id,
|
||||
message_id=client_message_id,
|
||||
)
|
||||
|
||||
assert response.user_message_id == client_message_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_setup_agent_message_has_message_id(
|
||||
self,
|
||||
acp_agent_loop: VibeAcpAgentLoop,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
env_file = tmp_path / ".env"
|
||||
|
||||
class FakeGlobalEnvFile:
|
||||
path = env_file
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.proxy_setup.GLOBAL_ENV_FILE", FakeGlobalEnvFile()
|
||||
)
|
||||
|
||||
session_response = await acp_agent_loop.new_session(
|
||||
cwd=str(Path.cwd()), mcp_servers=[]
|
||||
)
|
||||
_get_fake_client(acp_agent_loop)._session_updates.clear()
|
||||
|
||||
await acp_agent_loop.prompt(
|
||||
prompt=[TextContentBlock(type="text", text="/proxy-setup")],
|
||||
session_id=session_response.session_id,
|
||||
)
|
||||
|
||||
message_updates = [
|
||||
u
|
||||
for u in _get_fake_client(acp_agent_loop)._session_updates
|
||||
if isinstance(u.update, AgentMessageChunk)
|
||||
]
|
||||
assert len(message_updates) == 1
|
||||
assert message_updates[0].update.message_id is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_setup_unsets_value(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -38,6 +38,12 @@ class MockWidget:
|
|||
return self._get_selection_result
|
||||
|
||||
|
||||
class MockWidgetNoScreen:
|
||||
@property
|
||||
def text_selection(self) -> object:
|
||||
raise RuntimeError("node has no screen")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_app() -> App:
|
||||
app = MagicMock(spec=App)
|
||||
|
|
@ -73,6 +79,7 @@ def mock_app() -> App:
|
|||
],
|
||||
"empty text",
|
||||
),
|
||||
([MockWidgetNoScreen()], "widget with no screen (text_selection raises)"),
|
||||
],
|
||||
)
|
||||
def test_copy_selection_to_clipboard_no_notification(
|
||||
|
|
@ -87,6 +94,22 @@ def test_copy_selection_to_clipboard_no_notification(
|
|||
mock_app.notify.assert_not_called()
|
||||
|
||||
|
||||
@patch("vibe.cli.clipboard._copy_to_clipboard")
|
||||
def test_copy_selection_skips_detached_widget_and_collects_valid(
|
||||
mock_copy_to_clipboard: MagicMock, mock_app: MagicMock
|
||||
) -> None:
|
||||
detached = MockWidgetNoScreen()
|
||||
valid = MockWidget(
|
||||
text_selection=SimpleNamespace(), get_selection_result=("valid text", None)
|
||||
)
|
||||
mock_app.query.return_value = [detached, valid]
|
||||
|
||||
result = copy_selection_to_clipboard(mock_app)
|
||||
|
||||
assert result == "valid text"
|
||||
mock_copy_to_clipboard.assert_called_once_with("valid text")
|
||||
|
||||
|
||||
@patch("vibe.cli.clipboard._copy_to_clipboard")
|
||||
def test_copy_selection_to_clipboard_success(
|
||||
mock_copy_to_clipboard: MagicMock, mock_app: MagicMock
|
||||
|
|
|
|||
74
tests/core/test_backend_error.py
Normal file
74
tests/core/test_backend_error.py
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.llm.exceptions import BackendError, PayloadSummary
|
||||
|
||||
|
||||
def _make_payload_summary() -> PayloadSummary:
|
||||
return PayloadSummary(
|
||||
model="test-model",
|
||||
message_count=1,
|
||||
approx_chars=10,
|
||||
temperature=0.7,
|
||||
has_tools=False,
|
||||
tool_choice=None,
|
||||
)
|
||||
|
||||
|
||||
def _make_error(
|
||||
*, status: int | None, headers: dict[str, str] | None = None
|
||||
) -> BackendError:
|
||||
return BackendError(
|
||||
provider="test-provider",
|
||||
endpoint="/v1/chat/completions",
|
||||
status=status,
|
||||
reason="some reason",
|
||||
headers=headers or {},
|
||||
body_text="body",
|
||||
parsed_error=None,
|
||||
model="test-model",
|
||||
payload_summary=_make_payload_summary(),
|
||||
)
|
||||
|
||||
|
||||
class TestBackendErrorFmt:
|
||||
def test_standard_status_code(self) -> None:
|
||||
err = _make_error(status=500)
|
||||
msg = str(err)
|
||||
assert "500 Internal Server Error" in msg
|
||||
assert "test-provider" in msg
|
||||
|
||||
def test_non_standard_status_code(self) -> None:
|
||||
"""Status 529 is not in HTTPStatus and previously raised ValueError."""
|
||||
err = _make_error(status=529)
|
||||
msg = str(err)
|
||||
assert "529" in msg
|
||||
# Should not contain a phrase since 529 is not standard
|
||||
assert "LLM backend error [test-provider]" in msg
|
||||
|
||||
def test_no_status(self) -> None:
|
||||
err = _make_error(status=None)
|
||||
msg = str(err)
|
||||
assert "status: N/A" in msg
|
||||
|
||||
def test_unauthorized_short_circuits(self) -> None:
|
||||
err = _make_error(status=401)
|
||||
assert str(err) == "Invalid API key. Please check your API key and try again."
|
||||
|
||||
def test_rate_limit_short_circuits(self) -> None:
|
||||
err = _make_error(status=429)
|
||||
assert (
|
||||
str(err) == "Rate limit exceeded. Please wait a moment before trying again."
|
||||
)
|
||||
|
||||
def test_request_id_from_headers(self) -> None:
|
||||
err = _make_error(status=500, headers={"x-request-id": "req-123"})
|
||||
assert "req-123" in str(err)
|
||||
|
||||
@pytest.mark.parametrize("code", [530, 599, 999])
|
||||
def test_other_non_standard_codes(self, code: int) -> None:
|
||||
err = _make_error(status=code)
|
||||
msg = str(err)
|
||||
assert str(code) in msg
|
||||
assert "LLM backend error" in msg
|
||||
31
tests/core/test_retry.py
Normal file
31
tests/core/test_retry.py
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from vibe.core.utils.retry import _is_retryable_http_error
|
||||
|
||||
|
||||
def _make_http_status_error(status_code: int) -> httpx.HTTPStatusError:
|
||||
response = httpx.Response(
|
||||
status_code=status_code, request=httpx.Request("GET", "https://example.com")
|
||||
)
|
||||
return httpx.HTTPStatusError(
|
||||
message=f"Error {status_code}", request=response.request, response=response
|
||||
)
|
||||
|
||||
|
||||
class TestIsRetryableHttpError:
|
||||
@pytest.mark.parametrize("code", [408, 409, 425, 429, 500, 502, 503, 504, 529])
|
||||
def test_retryable_codes(self, code: int) -> None:
|
||||
assert _is_retryable_http_error(_make_http_status_error(code)) is True
|
||||
|
||||
@pytest.mark.parametrize("code", [400, 401, 403, 404, 422])
|
||||
def test_non_retryable_codes(self, code: int) -> None:
|
||||
assert _is_retryable_http_error(_make_http_status_error(code)) is False
|
||||
|
||||
def test_non_http_error_returns_false(self) -> None:
|
||||
assert _is_retryable_http_error(ValueError("not http")) is False
|
||||
|
||||
def test_generic_exception_returns_false(self) -> None:
|
||||
assert _is_retryable_http_error(RuntimeError("boom")) is False
|
||||
|
|
@ -6,7 +6,7 @@ from acp import (
|
|||
Agent as AcpAgent,
|
||||
Client,
|
||||
CreateTerminalResponse,
|
||||
KillTerminalCommandResponse,
|
||||
KillTerminalResponse,
|
||||
ReadTextFileResponse,
|
||||
ReleaseTerminalResponse,
|
||||
RequestPermissionResponse,
|
||||
|
|
@ -112,7 +112,7 @@ class FakeClient(Client):
|
|||
|
||||
async def kill_terminal(
|
||||
self, session_id: str, terminal_id: str, **kwargs: Any
|
||||
) -> KillTerminalCommandResponse | None:
|
||||
) -> KillTerminalResponse | None:
|
||||
raise NotImplementedError()
|
||||
|
||||
async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
|
||||
|
|
|
|||
|
|
@ -516,26 +516,6 @@ async def test_fill_missing_tool_responses_inserts_placeholders() -> None:
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_assistant_after_tool_appends_understood() -> None:
|
||||
agent_loop = build_test_agent_loop(
|
||||
config=make_config(),
|
||||
agent_name=BuiltinAgentName.AUTO_APPROVE,
|
||||
backend=FakeBackend(mock_llm_chunk(content="ok")),
|
||||
)
|
||||
tool_msg = LLMMessage(
|
||||
role=Role.tool, tool_call_id="tc_z", name="todo", content="Done"
|
||||
)
|
||||
agent_loop.messages.reset([agent_loop.messages[0], tool_msg])
|
||||
|
||||
await act_and_collect_events(agent_loop, "Next")
|
||||
|
||||
# find the seeded tool message and ensure the next message is "Understood."
|
||||
idx = next(i for i, m in enumerate(agent_loop.messages) if m.role == Role.tool)
|
||||
assert agent_loop.messages[idx + 1].role == Role.assistant
|
||||
assert agent_loop.messages[idx + 1].content == "Understood."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_tool_calls_produce_correct_events(
|
||||
telemetry_events: list[dict],
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue