2.0.0

Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai>
Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai>
Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai>
Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai>
Co-Authored-By: Clément Siriex <clement.sirieix@mistral.ai>
Co-Authored-By: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai>
Co-Authored-By: Thaddee Tyl <thaddee.tyl@gmail.com>
Co-Authored-By: David Brochart <david.brochart@gmail.com>
Co-Authored-By: Joseph Guhlin <joseph.guhlin@gmail.com>
Co-Authored-By: Thomas Kenbeek <thomaskenbeek@gmail.com>
Co-Authored-By: Remenby31 <baptiste.cruvellier31@gmail.com>

tests/session/test_session_loader.py

@@ -0,0 +1,539 @@
+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!"