from __future__ import annotations from pathlib import Path from unittest.mock import AsyncMock from acp import RequestError from acp.schema import ( AgentMessageChunk, AgentThoughtChunk, ClientCapabilities, ToolCallProgress, ToolCallStart, UserMessageChunk, ) import pytest 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 WORKSPACE_TRUST_CAPABILITY, VibeAcpAgentLoop from vibe.core.agent_loop import AgentLoop from vibe.core.agents.models import BuiltinAgentName from vibe.core.config import ModelConfig, SessionLoggingConfig from vibe.core.trusted_folders import trusted_folders_manager from vibe.core.types import Role @pytest.fixture def acp_agent_with_session_config( backend: FakeBackend, temp_session_dir: Path, monkeypatch: pytest.MonkeyPatch ) -> tuple[VibeAcpAgentLoop, FakeClient]: session_config = SessionLoggingConfig( save_dir=str(temp_session_dir), session_prefix="session", enabled=True ) config = build_test_vibe_config( active_model="devstral-latest", models=[ ModelConfig( name="devstral-latest", provider="mistral", alias="devstral-latest" ), ModelConfig( name="devstral-small", provider="mistral", alias="devstral-small" ), ], 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() monkeypatch.setattr("vibe.acp.acp_agent_loop.AgentLoop", PatchedAgentLoop) monkeypatch.setattr(VibeAcpAgentLoop, "_load_config", lambda self: config) vibe_acp_agent = VibeAcpAgentLoop() client = FakeClient() vibe_acp_agent.on_connect(client) client.on_connect(vibe_acp_agent) return vibe_acp_agent, client class TestLoadSession: @pytest.mark.asyncio async def test_load_session_response_structure( 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 = "test-sess-12345678" cwd = str(Path.cwd()) create_test_session(temp_session_dir, session_id, cwd) response = await acp_agent.load_session( cwd=cwd, mcp_servers=[], session_id=session_id ) assert response is not None assert response.models is not None assert len(response.models.available_models) == 2 assert response.models.current_model_id == "devstral-latest" assert response.models.available_models[0].model_id == "devstral-latest" assert response.models.available_models[0].name == "devstral-latest" assert response.models.available_models[1].model_id == "devstral-small" assert response.models.available_models[1].name == "devstral-small" assert response.modes is not None assert response.modes.current_mode_id == BuiltinAgentName.DEFAULT modes_ids = {m.id for m in response.modes.available_modes} assert modes_ids == { BuiltinAgentName.DEFAULT, BuiltinAgentName.CHAT, BuiltinAgentName.AUTO_APPROVE, BuiltinAgentName.PLAN, BuiltinAgentName.ACCEPT_EDITS, } assert response.config_options is not None assert len(response.config_options) == 3 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, BuiltinAgentName.AUTO_APPROVE, BuiltinAgentName.PLAN, BuiltinAgentName.ACCEPT_EDITS, } 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"} assert response.config_options[2].id == "thinking" assert response.config_options[2].category == "thinking" assert response.config_options[2].current_value == "off" @pytest.mark.asyncio async def test_load_session_resolves_workspace_trust_before_loading_config( self, acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient], temp_session_dir: Path, create_test_session, tmp_working_directory: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: acp_agent, client = acp_agent_with_session_config acp_agent.client_capabilities = ClientCapabilities( field_meta={WORKSPACE_TRUST_CAPABILITY: True} ) request_trust = AsyncMock(return_value={"decision": "trust_cwd"}) monkeypatch.setattr(client, "ext_method", request_trust) (tmp_working_directory / "AGENTS.md").write_text( "Loaded session project instructions", encoding="utf-8" ) session_id = "test-sess-trust" create_test_session(temp_session_dir, session_id, str(tmp_working_directory)) await acp_agent.load_session( cwd=str(tmp_working_directory), mcp_servers=[], session_id=session_id ) request_trust.assert_awaited_once() assert trusted_folders_manager.is_trusted(tmp_working_directory) is True system_prompt = acp_agent.sessions[session_id].agent_loop.messages[0].content assert system_prompt is not None assert "Loaded session project instructions" in system_prompt @pytest.mark.asyncio async def test_load_session_registers_session_with_original_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 = "orig-id-12345678" cwd = str(Path.cwd()) create_test_session(temp_session_dir, session_id, cwd) await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id) assert session_id in acp_agent.sessions assert acp_agent.sessions[session_id].id == session_id assert acp_agent.sessions[session_id].agent_loop.session_id == session_id @pytest.mark.asyncio async def test_load_session_injects_messages_into_agent_loop( 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-test-12345678" cwd = str(Path.cwd()) messages = [ {"role": "system", "content": "You are helpful"}, {"role": "user", "content": "First question"}, {"role": "assistant", "content": "First answer"}, {"role": "user", "content": "Second question"}, {"role": "assistant", "content": "Second answer"}, ] create_test_session(temp_session_dir, session_id, cwd, messages=messages) await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id) session = acp_agent.sessions[session_id] non_system = [m for m in session.agent_loop.messages if m.role != Role.system] assert len(non_system) == 4 @pytest.mark.asyncio async def test_load_session_replays_user_messages( 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 = "replay-usr-123456" cwd = str(Path.cwd()) messages = [{"role": "user", "content": "Hello world"}] 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.content.text == "Hello world" @pytest.mark.asyncio async def test_load_session_replays_assistant_messages( 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 = "replay-ast-123456" cwd = str(Path.cwd()) messages = [ {"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello! How can I help?"}, ] 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.content.text == "Hello! How can I help?" @pytest.mark.asyncio async def test_load_session_replays_tool_calls( 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 = "replay-tool-12345" cwd = str(Path.cwd()) messages = [ {"role": "user", "content": "Read the file"}, { "role": "assistant", "content": "", "tool_calls": [ { "id": "call_123", "type": "function", "function": { "name": "read", "arguments": '{"file_path": "/tmp/test.txt"}', }, } ], }, {"role": "tool", "tool_call_id": "call_123", "content": "file contents"}, ] create_test_session(temp_session_dir, session_id, cwd, messages=messages) await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id) tool_call_starts = [ u for u in client._session_updates if isinstance(u.update, ToolCallStart) ] assert len(tool_call_starts) == 1 assert tool_call_starts[0].update.title == "read" assert tool_call_starts[0].update.tool_call_id == "call_123" tool_results = [ u for u in client._session_updates if isinstance(u.update, ToolCallProgress) ] assert len(tool_results) == 1 assert tool_results[0].update.tool_call_id == "call_123" assert tool_results[0].update.status == "completed" @pytest.mark.asyncio async def test_load_session_replays_reasoning_content( 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 = "replay-reason-123" cwd = str(Path.cwd()) messages = [ {"role": "user", "content": "Think about this"}, { "role": "assistant", "content": "Here is my answer", "reasoning_content": "Let me think step by step...", }, ] create_test_session(temp_session_dir, session_id, cwd, messages=messages) await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id) thought_updates = [ u for u in client._session_updates if isinstance(u.update, AgentThoughtChunk) ] assert len(thought_updates) == 1 assert thought_updates[0].update.content.text == "Let me think step by step..." @pytest.mark.asyncio async def test_load_session_replays_reasoning_before_assistant_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 = "replay-order-1234" cwd = str(Path.cwd()) messages = [ {"role": "user", "content": "Think about this"}, { "role": "assistant", "content": "Here is my answer", "reasoning_content": "Let me think step by step...", }, ] create_test_session(temp_session_dir, session_id, cwd, messages=messages) await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id) response_updates = [ update.update for update in client._session_updates if isinstance(update.update, (AgentThoughtChunk, AgentMessageChunk)) ] assert [type(update) for update in response_updates] == [ AgentThoughtChunk, AgentMessageChunk, ] assert response_updates[0].content.text == "Let me think step by step..." assert response_updates[1].content.text == "Here is my answer" @pytest.mark.asyncio async def test_load_session_not_found_raises_error( self, acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient] ) -> None: acp_agent, _client = acp_agent_with_session_config with pytest.raises(RequestError): await acp_agent.load_session( cwd=str(Path.cwd()), mcp_servers=[], session_id="nonexistent-session" ) @pytest.mark.asyncio async def test_load_session_replays_full_conversation( 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 = "full-conv-1234567" cwd = str(Path.cwd()) messages = [ {"role": "user", "content": "First message"}, {"role": "assistant", "content": "First response"}, {"role": "user", "content": "Second message"}, {"role": "assistant", "content": "Second response"}, ] 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) ] agent_updates = [ u for u in client._session_updates if isinstance(u.update, AgentMessageChunk) ] assert len(user_updates) == 2 assert len(agent_updates) == 2 assert user_updates[0].update.content.text == "First message" 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_load_session_restores_agent_loop_session_identity( 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 = "restore-id-12345678" parent_session_id = "parent-id-87654321" cwd = str(Path.cwd()) session_dir = create_test_session( temp_session_dir, session_id, cwd, parent_session_id=parent_session_id ) await acp_agent.load_session(cwd=cwd, mcp_servers=[], session_id=session_id) agent_loop = acp_agent.sessions[session_id].agent_loop assert agent_loop.session_id == session_id assert agent_loop.parent_session_id == parent_session_id assert agent_loop.session_logger.session_id == session_id assert agent_loop.session_logger.session_dir == session_dir assert agent_loop.session_logger.session_metadata is not None assert ( agent_loop.session_logger.session_metadata.parent_session_id == parent_session_id ) @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 )