"""Unit tests for SessionManager. Tests cover: - Session creation and retrieval - Conversation creation with prompts - Session concurrency semaphore """ from __future__ import annotations import pytest import asyncio from types import SimpleNamespace from unittest.mock import Mock from importlib import import_module import langbot_plugin.api.entities.builtin.provider.session as provider_session import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message import langbot_plugin.api.entities.builtin.provider.prompt as provider_prompt from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.pipeline.pool import ( ExecutionContextMismatchError, ExecutionContextRequiredError, ) TEST_CONTEXT = ExecutionContext( instance_uuid='instance-test', workspace_uuid='workspace-test', placement_generation=1, ) TEST_BOT_UUID = 'bot-123' def bind_query_context(query, *, context=TEST_CONTEXT, bot_uuid=TEST_BOT_UUID): """Attach the trusted runtime scope expected by the session manager.""" query.bot_uuid = bot_uuid query._execution_context = context return query def bind_session_context(session, query): """Make a mocked legacy Session belong to the Query execution scope.""" session.bot_uuid = query.bot_uuid session._langbot_session_key = ( TEST_CONTEXT.instance_uuid, TEST_CONTEXT.workspace_uuid, TEST_CONTEXT.placement_generation, query.bot_uuid, query.launcher_type.value, query.launcher_id, ) return session def scoped_query( *, workspace_uuid='workspace-test', bot_uuid=TEST_BOT_UUID, placement_generation=1, pipeline_uuid=None, ): """Create a small Query-like object with a complete trusted scope.""" context = ExecutionContext( instance_uuid='instance-test', workspace_uuid=workspace_uuid, placement_generation=placement_generation, bot_uuid=bot_uuid, pipeline_uuid=pipeline_uuid, query_uuid=f'query-{workspace_uuid}-{bot_uuid}', ) return SimpleNamespace( launcher_type=provider_session.LauncherTypes.PERSON, launcher_id='same-launcher', sender_id='same-sender', bot_uuid=bot_uuid, _execution_context=context, ) def get_session_module(): """Lazy import to avoid circular import issues.""" return import_module('langbot.pkg.provider.session.sessionmgr') class TestSessionManagerInit: """Tests for SessionManager initialization.""" def test_init_stores_app_reference(self): """Test that __init__ stores the Application reference.""" sessionmgr = get_session_module() mock_app = Mock() manager = sessionmgr.SessionManager(mock_app) assert manager.ap is mock_app def test_init_empty_session_list(self): """Test that session_list starts empty.""" sessionmgr = get_session_module() mock_app = Mock() manager = sessionmgr.SessionManager(mock_app) assert manager.session_list == [] @pytest.mark.asyncio async def test_initialize_empty(self): """Test that initialize does nothing (current implementation).""" sessionmgr = get_session_module() mock_app = Mock() manager = sessionmgr.SessionManager(mock_app) await manager.initialize() # Should not raise or change state assert manager.session_list == [] class TestSessionManagerGetSession: """Tests for get_session method.""" @pytest.fixture def mock_app_with_config(self): """Create mock app with instance config.""" mock_app = Mock() mock_app.instance_config = Mock() mock_app.instance_config.data = {'concurrency': {'session': 5}} return mock_app @pytest.fixture def sample_query(self): """Create sample query for testing.""" query = Mock(spec=pipeline_query.Query) query.launcher_type = provider_session.LauncherTypes.PERSON query.launcher_id = '12345' query.sender_id = '12345' return bind_query_context(query) @pytest.mark.asyncio async def test_creates_new_session_when_not_found(self, mock_app_with_config, sample_query): """Test that get_session creates new session when not found.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) session = await manager.get_session(sample_query) assert session is not None assert session.launcher_type == sample_query.launcher_type assert session.launcher_id == sample_query.launcher_id assert session.sender_id == sample_query.sender_id assert len(manager.session_list) == 1 @pytest.mark.asyncio async def test_returns_existing_session_when_found(self, mock_app_with_config, sample_query): """Test that get_session returns existing session when found.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) # First call creates session session1 = await manager.get_session(sample_query) # Second call should return same session session2 = await manager.get_session(sample_query) assert session1 is session2 assert len(manager.session_list) == 1 @pytest.mark.asyncio async def test_session_has_semaphore(self, mock_app_with_config, sample_query): """Test that created session has semaphore for concurrency.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) session = await manager.get_session(sample_query) assert hasattr(session, '_semaphore') assert session._semaphore is not None assert isinstance(session._semaphore, asyncio.Semaphore) @pytest.mark.asyncio async def test_different_launchers_have_different_sessions(self, mock_app_with_config): """Test that different launcher_id creates different sessions.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) query1 = Mock(spec=pipeline_query.Query) query1.launcher_type = provider_session.LauncherTypes.PERSON query1.launcher_id = 'user1' query1.sender_id = 'user1' bind_query_context(query1) query2 = Mock(spec=pipeline_query.Query) query2.launcher_type = provider_session.LauncherTypes.PERSON query2.launcher_id = 'user2' query2.sender_id = 'user2' bind_query_context(query2) session1 = await manager.get_session(query1) session2 = await manager.get_session(query2) assert session1 is not session2 assert len(manager.session_list) == 2 @pytest.mark.asyncio async def test_different_launcher_types_have_different_sessions(self, mock_app_with_config): """Test that different launcher_type creates different sessions.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) query1 = Mock(spec=pipeline_query.Query) query1.launcher_type = provider_session.LauncherTypes.PERSON query1.launcher_id = 'same_id' query1.sender_id = 'same_id' bind_query_context(query1) query2 = Mock(spec=pipeline_query.Query) query2.launcher_type = provider_session.LauncherTypes.GROUP query2.launcher_id = 'same_id' query2.sender_id = 'same_id' bind_query_context(query2) session1 = await manager.get_session(query1) session2 = await manager.get_session(query2) assert session1 is not session2 assert len(manager.session_list) == 2 class TestSessionManagerGetConversation: """Tests for get_conversation method.""" @pytest.fixture def mock_app_with_config(self): """Create mock app with instance config.""" mock_app = Mock() mock_app.instance_config = Mock() mock_app.instance_config.data = {'concurrency': {'session': 5}} return mock_app @pytest.fixture def sample_session(self): """Create sample session for testing.""" session = Mock(spec=provider_session.Session) session.launcher_type = provider_session.LauncherTypes.PERSON session.launcher_id = '12345' session.sender_id = '12345' session.conversations = [] session.using_conversation = None return session @pytest.fixture def sample_query(self): """Create sample query for testing.""" query = Mock(spec=pipeline_query.Query) query.launcher_type = provider_session.LauncherTypes.PERSON query.launcher_id = '12345' query.sender_id = '12345' return bind_query_context(query) @pytest.mark.asyncio async def test_creates_conversation_with_prompt(self, mock_app_with_config, sample_query, sample_session): """Test that get_conversation creates conversation with prompt.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) bind_session_context(sample_session, sample_query) prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}] pipeline_uuid = 'pipeline-123' bot_uuid = 'bot-123' conversation = await manager.get_conversation( sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid ) assert conversation is not None assert conversation.pipeline_uuid == pipeline_uuid assert conversation.bot_uuid == bot_uuid assert conversation.prompt is not None assert len(sample_session.conversations) == 1 @pytest.mark.asyncio async def test_uses_existing_conversation_when_pipeline_matches( self, mock_app_with_config, sample_query, sample_session ): """Test that get_conversation uses existing conversation when pipeline matches.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) bind_session_context(sample_session, sample_query) prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}] pipeline_uuid = 'pipeline-123' bot_uuid = 'bot-123' # First call creates conversation conv1 = await manager.get_conversation(sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid) # Second call with same pipeline should return same conversation conv2 = await manager.get_conversation(sample_query, sample_session, prompt_config, pipeline_uuid, bot_uuid) assert conv1 is conv2 assert len(sample_session.conversations) == 1 @pytest.mark.asyncio async def test_creates_new_conversation_when_pipeline_changes( self, mock_app_with_config, sample_query, sample_session ): """Test that get_conversation creates new conversation when pipeline changes.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) bind_session_context(sample_session, sample_query) prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}] # First call with pipeline1 conv1 = await manager.get_conversation(sample_query, sample_session, prompt_config, 'pipeline-1', TEST_BOT_UUID) # Second call with different pipeline should create new conversation conv2 = await manager.get_conversation(sample_query, sample_session, prompt_config, 'pipeline-2', TEST_BOT_UUID) assert conv1 is not conv2 assert len(sample_session.conversations) == 2 assert sample_session.using_conversation is conv2 @pytest.mark.asyncio async def test_conversation_has_empty_messages(self, mock_app_with_config, sample_query, sample_session): """Test that created conversation has empty messages list.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) bind_session_context(sample_session, sample_query) prompt_config = [{'role': 'system', 'content': 'You are a helpful assistant.'}] conversation = await manager.get_conversation( sample_query, sample_session, prompt_config, 'pipeline-123', 'bot-123' ) assert conversation.messages == [] @pytest.mark.asyncio async def test_prompt_messages_from_config(self, mock_app_with_config, sample_query, sample_session): """Test that prompt messages are created from prompt_config.""" sessionmgr = get_session_module() manager = sessionmgr.SessionManager(mock_app_with_config) bind_session_context(sample_session, sample_query) prompt_config = [{'role': 'system', 'content': 'System message'}, {'role': 'user', 'content': 'User message'}] conversation = await manager.get_conversation( sample_query, sample_session, prompt_config, 'pipeline-123', 'bot-123' ) assert conversation.prompt.name == 'default' assert len(conversation.prompt.messages) == 2 class TestSessionManagerWorkspaceIsolation: """Regression coverage for workspace, bot, and placement fencing.""" @staticmethod def manager(): mock_app = Mock() mock_app.instance_config.data = {'concurrency': {'session': 5}} return get_session_module().SessionManager(mock_app) @pytest.mark.asyncio async def test_get_session_requires_trusted_query_scope(self): query = SimpleNamespace( launcher_type=provider_session.LauncherTypes.PERSON, launcher_id='same-launcher', sender_id='same-sender', bot_uuid=TEST_BOT_UUID, ) with pytest.raises(ExecutionContextRequiredError): await self.manager().get_session(query) @pytest.mark.asyncio async def test_same_launcher_in_two_workspaces_does_not_share_session(self): manager = self.manager() first = await manager.get_session(scoped_query(workspace_uuid='workspace-a')) second = await manager.get_session(scoped_query(workspace_uuid='workspace-b')) assert first is not second assert len(manager.session_list) == 2 @pytest.mark.asyncio async def test_same_launcher_in_two_bots_does_not_share_session(self): manager = self.manager() first = await manager.get_session(scoped_query(bot_uuid='bot-a')) second = await manager.get_session(scoped_query(bot_uuid='bot-b')) assert first is not second @pytest.mark.asyncio async def test_new_placement_generation_does_not_reuse_old_session(self): manager = self.manager() first = await manager.get_session(scoped_query(placement_generation=1)) second = await manager.get_session(scoped_query(placement_generation=2)) assert first is not second @pytest.mark.asyncio async def test_conversation_rejects_session_from_another_workspace(self): manager = self.manager() query_a = scoped_query(workspace_uuid='workspace-a') query_b = scoped_query(workspace_uuid='workspace-b') session_a = await manager.get_session(query_a) with pytest.raises(ExecutionContextMismatchError): await manager.get_conversation(query_b, session_a, [], 'pipeline-1', TEST_BOT_UUID) @pytest.mark.asyncio async def test_conversation_rejects_substituted_bot_argument(self): manager = self.manager() query = scoped_query(bot_uuid='bot-a') session = await manager.get_session(query) with pytest.raises(ExecutionContextMismatchError): await manager.get_conversation(query, session, [], 'pipeline-1', 'bot-b') @pytest.mark.asyncio async def test_per_workspace_capacity_evicts_oldest_idle_session(self): manager = self.manager() manager.ap.instance_config.data['system'] = { 'session_retention': { 'max_entries': 10, 'max_entries_per_workspace': 2, } } queries = [] sessions = [] for index in range(3): query = scoped_query() query.launcher_id = f'launcher-{index}' queries.append(query) sessions.append(await manager.get_session(query)) assert len(manager.session_list) == 2 assert sessions[0] not in manager.session_list assert sessions[1:] == manager.session_list assert await manager.get_session(queries[-1]) is manager.session_list[-1] @pytest.mark.asyncio async def test_new_session_does_not_scan_other_workspace_sessions(self): manager = self.manager() manager.ap.instance_config.data['system'] = { 'session_retention': { 'max_entries': 600, 'max_entries_per_workspace': 2, } } for index in range(512): query = scoped_query(workspace_uuid=f'workspace-{index}') query.launcher_id = f'launcher-{index}' await manager.get_session(query) class NoGlobalIterationDict(dict): def __iter__(self): raise AssertionError('global session index iteration is forbidden') def items(self): raise AssertionError('global session index iteration is forbidden') def values(self): raise AssertionError('global session index iteration is forbidden') manager._session_index = NoGlobalIterationDict(manager._session_index) query = scoped_query(workspace_uuid='workspace-new') query.launcher_id = 'launcher-new' session = await manager.get_session(query) assert session.workspace_uuid == 'workspace-new' assert len(manager._session_index) == 513 @pytest.mark.asyncio async def test_stale_expiry_revision_does_not_evict_recent_session( self, monkeypatch, ): sessionmgr = get_session_module() manager = self.manager() manager.ap.instance_config.data['system'] = { 'session_retention': { 'max_entries': 10, 'max_entries_per_workspace': 10, 'idle_ttl_seconds': 1, } } clock = [0.0] monkeypatch.setattr(sessionmgr.time, 'monotonic', lambda: clock[0]) first_query = scoped_query() first_query.launcher_id = 'first' first = await manager.get_session(first_query) clock[0] = 0.5 assert await manager.get_session(first_query) is first clock[0] = 1.25 second_query = scoped_query() second_query.launcher_id = 'second' await manager.get_session(second_query) assert first in manager.session_list clock[0] = 2.0 third_query = scoped_query() third_query.launcher_id = 'third' await manager.get_session(third_query) assert first not in manager.session_list @pytest.mark.asyncio async def test_access_revision_heap_stays_bounded(self): manager = self.manager() query = scoped_query() await manager.get_session(query) for _ in range(1000): await manager.get_session(query) assert len(manager._session_expiry_heap) <= 64 def test_trim_conversation_drops_retained_binary_payloads(self): manager = self.manager() conversation = provider_session.Conversation( prompt=provider_prompt.Prompt(name='test', messages=[]), messages=[ provider_message.Message( role='user', content=[ provider_message.ContentElement.from_text('hello'), provider_message.ContentElement.from_image_base64('x' * 1000000), provider_message.ContentElement.from_file_base64( 'y' * 1000000, 'large.bin', ), ], ) ], pipeline_uuid='pipeline-1', bot_uuid=TEST_BOT_UUID, ) manager.trim_conversation_messages(conversation, max_rounds=10) content = conversation.messages[0].content assert content[1].image_base64 is None assert content[2].file_base64 is None