from __future__ import annotations from datetime import datetime import json from pathlib import Path import time import pytest from vibe.core.config import SessionLoggingConfig from vibe.core.session.session_loader import SessionLoader from vibe.core.types import LLMMessage, Role, ToolCall @pytest.fixture def temp_session_dir(tmp_path: Path) -> Path: """Create a temporary directory for session loader tests.""" session_dir = tmp_path / "sessions" session_dir.mkdir() return session_dir @pytest.fixture def session_config(temp_session_dir: Path) -> SessionLoggingConfig: """Create a session logging config for testing.""" return SessionLoggingConfig( save_dir=str(temp_session_dir), session_prefix="test", enabled=True ) @pytest.fixture def create_test_session(): """Helper fixture to create a test session with messages and metadata.""" def _create_test_session( session_dir: Path, session_id: str, messages: list[LLMMessage] | None = None, metadata: dict | None = None, ) -> Path: """Create a test session directory with messages and metadata files.""" # Create session directory timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") session_folder = session_dir / f"test_{timestamp}_{session_id[:8]}" session_folder.mkdir(exist_ok=True) # Create messages file messages_file = session_folder / "messages.jsonl" if messages is None: messages = [ LLMMessage(role=Role.system, content="System prompt"), LLMMessage(role=Role.user, content="Hello"), LLMMessage(role=Role.assistant, content="Hi there!"), ] with messages_file.open("w", encoding="utf-8") as f: for message in messages: f.write( json.dumps( message.model_dump(exclude_none=True), ensure_ascii=False ) + "\n" ) # Create metadata file metadata_file = session_folder / "meta.json" if metadata is None: metadata = { "session_id": session_id, "start_time": "2023-01-01T12:00:00", "end_time": "2023-01-01T12:05:00", "total_messages": 2, "stats": { "steps": 1, "session_prompt_tokens": 10, "session_completion_tokens": 20, }, "system_prompt": {"content": "System prompt", "role": "system"}, } with metadata_file.open("w", encoding="utf-8") as f: json.dump(metadata, f, indent=2) return session_folder return _create_test_session class TestSessionLoaderFindLatestSession: def test_find_latest_session_no_sessions( self, session_config: SessionLoggingConfig ) -> None: """Test finding latest session when no sessions exist.""" result = SessionLoader.find_latest_session(session_config) assert result is None def test_find_latest_session_single_session( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding latest session with a single session.""" session_dir = Path(session_config.save_dir) session = create_test_session(session_dir, "session-123") result = SessionLoader.find_latest_session(session_config) assert result is not None assert result.exists() assert result == session def test_find_latest_session_multiple_sessions( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding latest session with multiple sessions.""" session_dir = Path(session_config.save_dir) create_test_session(session_dir, "session-123") time.sleep(0.01) create_test_session(session_dir, "session-456") time.sleep(0.01) latest = create_test_session(session_dir, "session-789") result = SessionLoader.find_latest_session(session_config) assert result is not None assert result.exists() assert result == latest def test_find_latest_session_nonexistent_save_dir(self) -> None: """Test finding latest session when save directory doesn't exist.""" # Modify config to point to non-existent directory bad_config = SessionLoggingConfig( save_dir="/nonexistent/path", session_prefix="test", enabled=True ) result = SessionLoader.find_latest_session(bad_config) assert result is None def test_find_latest_session_with_invalid_sessions( self, session_config: SessionLoggingConfig ) -> None: """Test finding latest session when only invalid sessions exist.""" session_dir = Path(session_config.save_dir) invalid_session1 = session_dir / "test_20230101_120000_invalid1" invalid_session1.mkdir() (invalid_session1 / "messages.jsonl").write_text("[]") # Missing meta.json invalid_session2 = session_dir / "test_20230101_120001_invalid2" invalid_session2.mkdir() (invalid_session2 / "meta.json").write_text('{"session_id": "invalid"}') result = SessionLoader.find_latest_session(session_config) assert result is None # Should return None when no valid sessions exist def test_find_latest_session_with_mixed_valid_invalid( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding latest session when both valid and invalid sessions exist.""" session_dir = Path(session_config.save_dir) invalid_session = session_dir / "test_20230101_120000_invalid" invalid_session.mkdir() (invalid_session / "messages.jsonl").write_text( '{"role": "user", "content": "test"}\n' ) time.sleep(0.01) valid_session = create_test_session(session_dir, "test_20230101_120001_valid") time.sleep(0.01) newest_invalid = session_dir / "test_20230101_120002_newest" newest_invalid.mkdir() (newest_invalid / "messages.jsonl").write_text( '{"role": "user", "content": "test"}\n' ) result = SessionLoader.find_latest_session(session_config) assert result is not None assert result == valid_session def test_find_latest_session_with_invalid_json( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding latest session when sessions have invalid JSON.""" session_dir = Path(session_config.save_dir) invalid_meta_session = session_dir / "test_20230101_120000_invalid_meta" invalid_meta_session.mkdir() (invalid_meta_session / "messages.jsonl").write_text( '{"role": "user", "content": "test"}\n' ) (invalid_meta_session / "meta.json").write_text("{invalid json}") time.sleep(0.01) invalid_msg_session = session_dir / "test_20230101_120001_invalid_msg" invalid_msg_session.mkdir() (invalid_msg_session / "messages.jsonl").write_text("{invalid json}") (invalid_msg_session / "meta.json").write_text('{"session_id": "invalid"}') time.sleep(0.01) valid_session = create_test_session(session_dir, "test_20230101_120002_valid") result = SessionLoader.find_latest_session(session_config) assert result is not None assert result == valid_session class TestSessionLoaderFindSessionById: def test_find_session_by_id_exact_match( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding session by exact ID match.""" session_dir = Path(session_config.save_dir) session_folder = create_test_session(session_dir, "test-session-123") # Test with full UUID format result = SessionLoader.find_session_by_id("test-session-123", session_config) assert result is not None assert result == session_folder def test_find_session_by_id_short_uuid( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding session by short UUID.""" session_dir = Path(session_config.save_dir) session_folder = create_test_session( session_dir, "abc12345-6789-0123-4567-89abcdef0123" ) # Test with short UUID result = SessionLoader.find_session_by_id("abc12345", session_config) assert result is not None assert result == session_folder def test_find_session_by_id_partial_match( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding session by partial ID match""" session_dir = Path(session_config.save_dir) session_folder = create_test_session(session_dir, "abc12345678") # Test with partial match result = SessionLoader.find_session_by_id("abc12345", session_config) assert result is not None assert result == session_folder def test_find_session_by_id_multiple_matches( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding session when multiple sessions match (should return most recent).""" session_dir = Path(session_config.save_dir) # Create first session create_test_session(session_dir, "abcd1234") # Sleep to ensure different modification times time.sleep(0.01) # Create second session with similar ID prefix session_2 = create_test_session(session_dir, "abcd1234") result = SessionLoader.find_session_by_id("abcd1234", session_config) assert result is not None assert result == session_2 def test_find_session_by_id_no_match( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding session by ID when no match exists.""" session_dir = Path(session_config.save_dir) create_test_session(session_dir, "test-session-123") result = SessionLoader.find_session_by_id("nonexistent", session_config) assert result is None def test_find_session_by_id_nonexistent_save_dir(self) -> None: """Test finding session by ID when save directory doesn't exist.""" bad_config = SessionLoggingConfig( save_dir="/nonexistent/path", session_prefix="test", enabled=True ) result = SessionLoader.find_session_by_id("test-session", bad_config) assert result is None class TestSessionLoaderLoadSession: def test_load_session_success( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test successfully loading a session.""" session_dir = Path(session_config.save_dir) session_folder = create_test_session(session_dir, "test-session-123") messages, metadata = SessionLoader.load_session(session_folder) # Verify messages assert len(messages) == 2 assert messages[0].role == "user" assert messages[0].content == "Hello" assert messages[1].role == "assistant" assert messages[1].content == "Hi there!" # Verify metadata assert metadata["session_id"] == "test-session-123" assert metadata["total_messages"] == 2 assert "stats" in metadata assert "system_prompt" in metadata def test_load_session_empty_messages( self, session_config: SessionLoggingConfig ) -> None: """Test loading session with empty messages file.""" session_dir = Path(session_config.save_dir) session_folder = session_dir / "test_20230101_120000_test123" session_folder.mkdir() # Create empty messages file messages_file = session_folder / "messages.jsonl" messages_file.write_text("") # Create metadata file metadata_file = session_folder / "meta.json" metadata_file.write_text('{"session_id": "test-session"}') with pytest.raises(ValueError, match="Session messages file is empty"): SessionLoader.load_session(session_folder) def test_load_session_invalid_json_messages( self, session_config: SessionLoggingConfig ) -> None: """Test loading session with invalid JSON in messages file.""" session_dir = Path(session_config.save_dir) session_folder = session_dir / "test_20230101_120000_test123" session_folder.mkdir() # Create messages file with invalid JSON messages_file = session_folder / "messages.jsonl" messages_file.write_text("{invalid json}") # Create metadata file metadata_file = session_folder / "meta.json" metadata_file.write_text('{"session_id": "test-session"}') with pytest.raises(ValueError, match="Session messages contain invalid JSON"): SessionLoader.load_session(session_folder) def test_load_session_invalid_json_metadata( self, session_config: SessionLoggingConfig ) -> None: """Test loading session with invalid JSON in metadata file.""" session_dir = Path(session_config.save_dir) session_folder = session_dir / "test_20230101_120000_test123" session_folder.mkdir() # Create valid messages file messages_file = session_folder / "messages.jsonl" messages_file.write_text('{"role": "user", "content": "Hello"}\n') # Create metadata file with invalid JSON metadata_file = session_folder / "meta.json" metadata_file.write_text("{invalid json}") with pytest.raises(ValueError, match="Session metadata contains invalid JSON"): SessionLoader.load_session(session_folder) def test_load_session_no_metadata_file( self, session_config: SessionLoggingConfig ) -> None: """Test loading session when metadata file doesn't exist.""" session_dir = Path(session_config.save_dir) session_folder = session_dir / "test_20230101_120000_test123" session_folder.mkdir() # Create valid messages file using the same format as create_test_session messages = [ LLMMessage(role=Role.system, content="System prompt"), LLMMessage(role=Role.user, content="Hello"), LLMMessage(role=Role.assistant, content="Hi there!"), ] messages_file = session_folder / "messages.jsonl" with messages_file.open("w", encoding="utf-8") as f: for message in messages: f.write( json.dumps( message.model_dump(exclude_none=True), ensure_ascii=False ) + "\n" ) loaded_messages, metadata = SessionLoader.load_session(session_folder) assert len(loaded_messages) == 2 assert loaded_messages[0].content == "Hello" assert loaded_messages[0].role == Role.user assert loaded_messages[1].content == "Hi there!" assert loaded_messages[1].role == Role.assistant assert metadata == {} def test_load_session_nonexistent_directory( self, session_config: SessionLoggingConfig ) -> None: """Test loading session from non-existent directory.""" nonexistent_dir = Path(session_config.save_dir) / "nonexistent" with pytest.raises(ValueError, match="Error reading session messages"): SessionLoader.load_session(nonexistent_dir) class TestSessionLoaderEdgeCases: def test_find_latest_session_with_different_prefixes( self, session_config: SessionLoggingConfig ) -> None: """Test finding latest session when sessions have different prefixes.""" session_dir = Path(session_config.save_dir) # Create sessions with different prefixes other_session = session_dir / "other_20230101_120000_test123" other_session.mkdir() (other_session / "messages.jsonl").write_text( '{"role": "user", "content": "test"}\n' ) test_session = session_dir / "test_20230101_120000_test456" test_session.mkdir() (test_session / "messages.jsonl").write_text( '{"role": "user", "content": "test"}\n' ) (test_session / "meta.json").write_text('{"session_id": "test456"}') result = SessionLoader.find_latest_session(session_config) assert result is not None assert result.name == "test_20230101_120000_test456" def test_find_session_by_id_with_special_characters( self, session_config: SessionLoggingConfig, create_test_session ) -> None: """Test finding session by ID containing special characters.""" session_dir = Path(session_config.save_dir) session_folder = create_test_session( session_dir, "test-session_with-special.chars" ) result = SessionLoader.find_session_by_id( "test-session_with-special.chars", session_config ) assert result is not None assert result == session_folder def test_load_session_with_complex_messages( self, session_config: SessionLoggingConfig ) -> None: """Test loading session with complex message structures.""" session_dir = Path(session_config.save_dir) session_folder = session_dir / "test_20230101_120000_test123" session_folder.mkdir() # Create messages with complex structure complex_messages = [ LLMMessage(role=Role.system, content="System prompt"), LLMMessage( role=Role.user, content="Complex message", reasoning_content="Some reasoning", tool_calls=[ToolCall(id="call1", index=1, type="function")], ), LLMMessage( role=Role.assistant, content="Response", tool_calls=[ToolCall(id="call2", index=2, type="function")], ), ] messages_file = session_folder / "messages.jsonl" with messages_file.open("w", encoding="utf-8") as f: for message in complex_messages: f.write( json.dumps( message.model_dump(exclude_none=True), ensure_ascii=False ) + "\n" ) # Create metadata file metadata_file = session_folder / "meta.json" metadata_file.write_text('{"session_id": "complex-session"}') messages, _ = SessionLoader.load_session(session_folder) # Verify complex messages are loaded correctly assert len(messages) == 2 assert messages[0].role == Role.user assert messages[0].content == "Complex message" assert messages[0].reasoning_content == "Some reasoning" assert len(messages[0].tool_calls or []) == 1 assert messages[1].role == Role.assistant assert len(messages[1].tool_calls or []) == 1 assert messages[1].content == "Response" def test_load_session_system_prompt_ignored_in_messages( self, session_config: SessionLoggingConfig ) -> None: """Test that system prompt is ignored when written in messages.jsonl.""" session_dir = Path(session_config.save_dir) session_folder = session_dir / "test_20230101_120000_test123" session_folder.mkdir() messages_with_system = [ LLMMessage(role=Role.system, content="System prompt from messages"), LLMMessage(role=Role.user, content="Hello"), LLMMessage(role=Role.assistant, content="Hi there!"), ] messages_file = session_folder / "messages.jsonl" with messages_file.open("w", encoding="utf-8") as f: for message in messages_with_system: f.write( json.dumps( message.model_dump(exclude_none=True), ensure_ascii=False ) + "\n" ) metadata_file = session_folder / "meta.json" metadata_file.write_text( json.dumps({"session_id": "test-session", "total_messages": 3}) ) messages, metadata = SessionLoader.load_session(session_folder) # Verify that system prompt from messages.jsonl is ignored assert len(messages) == 2 assert messages[0].role == Role.user assert messages[0].content == "Hello" assert messages[1].role == Role.assistant assert messages[1].content == "Hi there!"