b23f49e5f4
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: allansimon-mistral <allan.simon@ext.mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
117 lines
4.4 KiB
Python
117 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
from acp import PROTOCOL_VERSION, RequestError
|
|
from acp.schema import TextContentBlock
|
|
import pytest
|
|
from pytest import raises
|
|
|
|
from tests.mock.utils import mock_llm_chunk
|
|
from tests.stubs.fake_backend import FakeBackend
|
|
from vibe.acp.acp_agent_loop import VibeAcpAgentLoop
|
|
from vibe.core.types import Role
|
|
|
|
|
|
class TestMultiSessionCore:
|
|
@pytest.mark.asyncio
|
|
async def test_different_sessions_use_different_agents(
|
|
self, acp_agent_loop: VibeAcpAgentLoop
|
|
) -> None:
|
|
await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION)
|
|
session1_response = await acp_agent_loop.new_session(
|
|
cwd=str(Path.cwd()), mcp_servers=[]
|
|
)
|
|
session1 = acp_agent_loop.sessions[session1_response.session_id]
|
|
session2_response = await acp_agent_loop.new_session(
|
|
cwd=str(Path.cwd()), mcp_servers=[]
|
|
)
|
|
session2 = acp_agent_loop.sessions[session2_response.session_id]
|
|
|
|
assert session1.id != session2.id
|
|
# Each agent loop should be independent
|
|
assert session1.agent_loop is not session2.agent_loop
|
|
assert id(session1.agent_loop) != id(session2.agent_loop)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_error_on_nonexistent_session(
|
|
self, acp_agent_loop: VibeAcpAgentLoop
|
|
) -> None:
|
|
await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION)
|
|
await acp_agent_loop.new_session(cwd=str(Path.cwd()), mcp_servers=[])
|
|
|
|
fake_session_id = "fake-session-id-" + str(uuid4())
|
|
|
|
with raises(RequestError) as exc_info:
|
|
await acp_agent_loop.prompt(
|
|
session_id=fake_session_id,
|
|
prompt=[TextContentBlock(type="text", text="Hello, world!")],
|
|
)
|
|
|
|
assert isinstance(exc_info.value, RequestError)
|
|
assert exc_info.value.code == -32602
|
|
assert "Session not found" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_simultaneous_message_processing(
|
|
self, acp_agent_loop: VibeAcpAgentLoop, backend: FakeBackend
|
|
) -> None:
|
|
await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION)
|
|
session1_response = await acp_agent_loop.new_session(
|
|
cwd=str(Path.cwd()), mcp_servers=[]
|
|
)
|
|
session1 = acp_agent_loop.sessions[session1_response.session_id]
|
|
session2_response = await acp_agent_loop.new_session(
|
|
cwd=str(Path.cwd()), mcp_servers=[]
|
|
)
|
|
session2 = acp_agent_loop.sessions[session2_response.session_id]
|
|
|
|
backend._streams = [
|
|
[mock_llm_chunk(content="Response 1")],
|
|
[mock_llm_chunk(content="Response 2")],
|
|
]
|
|
|
|
async def run_session1():
|
|
await acp_agent_loop.prompt(
|
|
session_id=session1.id,
|
|
prompt=[TextContentBlock(type="text", text="Prompt for session 1")],
|
|
)
|
|
|
|
async def run_session2():
|
|
await acp_agent_loop.prompt(
|
|
session_id=session2.id,
|
|
prompt=[TextContentBlock(type="text", text="Prompt for session 2")],
|
|
)
|
|
|
|
await asyncio.gather(run_session1(), run_session2())
|
|
|
|
user_message1 = next(
|
|
(msg for msg in session1.agent_loop.messages if msg.role == Role.user), None
|
|
)
|
|
assert user_message1 is not None
|
|
assert user_message1.content == "Prompt for session 1"
|
|
user_message2 = next(
|
|
(msg for msg in session2.agent_loop.messages if msg.role == Role.user), None
|
|
)
|
|
assert user_message2 is not None
|
|
assert user_message2.content == "Prompt for session 2"
|
|
|
|
# Backend stream order is non-deterministic under asyncio.gather, so
|
|
# assert that both sessions received distinct responses from the
|
|
# expected set rather than pinning a specific assignment.
|
|
assistant_message1 = next(
|
|
(msg for msg in session1.agent_loop.messages if msg.role == Role.assistant),
|
|
None,
|
|
)
|
|
assistant_message2 = next(
|
|
(msg for msg in session2.agent_loop.messages if msg.role == Role.assistant),
|
|
None,
|
|
)
|
|
assert assistant_message1 is not None
|
|
assert assistant_message2 is not None
|
|
assert {assistant_message1.content, assistant_message2.content} == {
|
|
"Response 1",
|
|
"Response 2",
|
|
}
|