from __future__ import annotations import asyncio from pathlib import Path from acp.schema import TextContentBlock import pytest from vibe.acp.acp_agent_loop import VibeAcpAgentLoop from vibe.acp.exceptions import InvalidRequestError from vibe.core.agents.models import BuiltinAgentName from vibe.core.session.session_id import extract_suffix from vibe.core.types import FunctionCall, LLMMessage, Role, ToolCall class TestACPForkSession: @pytest.mark.asyncio async def test_fork_session_clones_history_and_mode( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] await acp_agent_loop.set_session_mode( session_id=source_session.id, mode_id=BuiltinAgentName.PLAN ) await acp_agent_loop.prompt( prompt=[TextContentBlock(type="text", text="Say hi")], session_id=source_session.id, ) response = await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[] ) assert response.session_id != source_session.id assert response.modes is not None assert response.modes.current_mode_id == BuiltinAgentName.PLAN assert response.models is not None assert ( response.models.current_model_id == source_session.agent_loop.config.active_model ) assert response.config_options is not None assert len(response.config_options) == 3 forked_session = acp_agent_loop.sessions[response.session_id] assert forked_session.agent_loop.agent_profile.name == BuiltinAgentName.PLAN source_messages = [ message.model_dump(mode="json", exclude_none=True) for message in source_session.agent_loop.messages if message.role != Role.system ] forked_messages = [ message.model_dump(mode="json", exclude_none=True) for message in forked_session.agent_loop.messages if message.role != Role.system ] assert forked_messages == source_messages @pytest.mark.asyncio async def test_fork_session_from_user_message_keeps_full_turn( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] source_session.agent_loop.messages.extend([ LLMMessage(role=Role.user, content="First", message_id="user-1"), LLMMessage( role=Role.assistant, content="First answer", message_id="assistant-1" ), LLMMessage( role=Role.assistant, content="", message_id="assistant-2", tool_calls=[ ToolCall( id="call-1", index=0, function=FunctionCall( name="read", arguments='{"file_path":"a.txt"}' ), ) ], ), LLMMessage(role=Role.tool, content="contents", tool_call_id="call-1"), LLMMessage(role=Role.user, content="Second", message_id="user-2"), LLMMessage( role=Role.assistant, content="Second answer", message_id="assistant-3" ), ]) response = await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[], messageId="user-1", ) forked_session = acp_agent_loop.sessions[response.session_id] assert [ (message.role, message.content, message.message_id, message.tool_call_id) for message in forked_session.agent_loop.messages if message.role != Role.system ] == [ (Role.user, "First", "user-1", None), (Role.assistant, "First answer", "assistant-1", None), (Role.assistant, "", "assistant-2", None), (Role.tool, "contents", None, "call-1"), ] @pytest.mark.asyncio async def test_fork_session_rejects_non_user_message_id( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] source_session.agent_loop.messages.extend([ LLMMessage(role=Role.user, content="First", message_id="user-1"), LLMMessage( role=Role.assistant, content="First answer", message_id="assistant-1" ), ]) with pytest.raises( InvalidRequestError, match="Fork from message_id is only supported for user messages", ): await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[], messageId="assistant-1", ) @pytest.mark.asyncio async def test_fork_session_sets_parent_session_id( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] response = await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[] ) forked_session = acp_agent_loop.sessions[response.session_id] assert forked_session.agent_loop.parent_session_id == source_session.id @pytest.mark.asyncio async def test_fork_session_inherits_limits( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] await acp_agent_loop.set_config_option( session_id=source_session.id, config_id="max_turns", value="42" ) await acp_agent_loop.set_config_option( session_id=source_session.id, config_id="max_tokens", value="8192" ) source_session.agent_loop._max_price = 1.5 source_session.agent_loop._max_session_tokens = 100_000 response = await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[] ) forked = acp_agent_loop.sessions[response.session_id].agent_loop assert forked._max_turns == 42 assert forked._max_tokens == 8192 assert forked._max_price == 1.5 assert forked._max_session_tokens == 100_000 @pytest.mark.asyncio async def test_fork_session_rejects_running_session( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] source_session.set_prompt_task(asyncio.sleep(10)) with pytest.raises( InvalidRequestError, match="Cannot fork a session while the agent loop is running", ): await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[] ) await source_session.cancel_prompt() @pytest.mark.asyncio async def test_fork_session_preserves_session_id_suffix( self, acp_agent_loop: VibeAcpAgentLoop ) -> None: session_response = await acp_agent_loop.new_session( cwd=str(Path.cwd()), mcp_servers=[] ) source_session = acp_agent_loop.sessions[session_response.session_id] response = await acp_agent_loop.fork_session( cwd=str(Path.cwd()), session_id=source_session.id, mcp_servers=[] ) assert extract_suffix(response.session_id) == extract_suffix(source_session.id)