dirigent/crates/dirigent_archivist/tests/pagination_test.rs

//! Pagination tests for dirigent_archivist
//!
//! These tests verify the count_messages and get_messages_range functionality.

#[cfg(test)]
mod pagination_tests {
    use chrono::Utc;
    use dirigent_archivist::{
        backends::JsonlBackend, Archivist, MessageRecord, RegisterConnectorRequest,
        RegisterSessionRequest, Result,
    };
    use std::sync::Arc;
    use uuid::Uuid;

    /// Build a self-contained coordinator rooted at `archive_root`, backed by
    /// a single `JsonlBackend`. Avoids the shared `.archives.json` race that
    /// `new_with_single_archive` creates in the tempdir's parent.
    async fn mk_archivist(archive_root: std::path::PathBuf) -> Result<Archivist> {
        let backend = Arc::new(JsonlBackend::new(archive_root).await?);
        Archivist::from_single_backend("main".into(), backend).await
    }

    #[tokio::test]
    async fn test_pagination_count_and_range() -> Result<()> {
        let temp_dir = std::env::temp_dir().join(format!("archivist_test_{}", Uuid::now_v7()));
        let archivist = mk_archivist(temp_dir.clone()).await?;

        // Register connector
        let connector_req = RegisterConnectorRequest {
            r#type: "OpenCode".to_string(),
            title: "Test Connector".to_string(),
            client_native_id: "opencode@localhost:12225".to_string(),
            custom_uid: None,
            metadata: serde_json::json!({}),
            fingerprint: None,
        };
        let connector_response = archivist.register_connector(connector_req, None).await?;

        // Register session
        let session_req = RegisterSessionRequest {
            connector_uid: connector_response.connector_uid,
            native_session_id: "native-123".to_string(),
            title: Some("Pagination Test".to_string()),
            custom_scroll_id: None,
            metadata: serde_json::json!({}),
            completeness: Default::default(),
            parent_scroll_id: None,
            is_subagent: false,
            continuation: None,
            agent_id: None,
            subagent_type: None,
            spawning_tool_use_id: None,
        };
        let session_response = archivist.register_session(session_req, None).await?;
        let scroll_id = session_response.scroll_id;

        // Test empty session
        let count = archivist.count_messages(scroll_id, None).await?;
        assert_eq!(count, 0, "Empty session should have 0 messages");

        let range = archivist.get_messages_range(scroll_id, 0, 10, None).await?;
        assert_eq!(range.len(), 0, "Empty session should return empty range");

        // Add 25 messages with varying timestamps
        let mut messages = Vec::new();
        let base_time = Utc::now();
        for i in 0..25 {
            messages.push(MessageRecord {
                version: 1,
                message_id: Uuid::now_v7(),
                session: scroll_id,
                parent_id: None,
                ts: base_time + chrono::Duration::seconds(i),
                role: if i % 2 == 0 { "user" } else { "assistant" }.to_string(),
                author: Some("test".to_string()),
                content_md: format!("Message {}", i),
                content_parts: None,
                attachments: Vec::new(),
                metadata: serde_json::json!({}),
            });
        }
        archivist.append_messages(scroll_id, messages, None).await?;

        // Test count_messages
        let count = archivist.count_messages(scroll_id, None).await?;
        assert_eq!(count, 25, "Should count 25 messages");

        // Test get_messages_range - first page
        let page1 = archivist.get_messages_range(scroll_id, 0, 10, None).await?;
        assert_eq!(page1.len(), 10, "First page should have 10 messages");
        assert_eq!(page1[0].content_md, "Message 0", "First message should be Message 0");
        assert_eq!(page1[9].content_md, "Message 9", "10th message should be Message 9");

        // Test get_messages_range - second page
        let page2 = archivist.get_messages_range(scroll_id, 10, 10, None).await?;
        assert_eq!(page2.len(), 10, "Second page should have 10 messages");
        assert_eq!(page2[0].content_md, "Message 10", "11th message should be Message 10");
        assert_eq!(page2[9].content_md, "Message 19", "20th message should be Message 19");

        // Test get_messages_range - partial last page
        let page3 = archivist.get_messages_range(scroll_id, 20, 10, None).await?;
        assert_eq!(page3.len(), 5, "Last page should have 5 messages");
        assert_eq!(page3[0].content_md, "Message 20", "21st message should be Message 20");
        assert_eq!(page3[4].content_md, "Message 24", "25th message should be Message 24");

        // Test get_messages_range - offset beyond messages
        let page4 = archivist.get_messages_range(scroll_id, 30, 10, None).await?;
        assert_eq!(page4.len(), 0, "Offset beyond messages should return empty");

        // Verify chronological ordering is maintained in pagination
        let all_messages = archivist.get_messages(scroll_id, None).await?;
        let first_10_from_all = &all_messages[0..10];
        let first_10_from_page = &page1[..];

        for i in 0..10 {
            assert_eq!(
                first_10_from_all[i].message_id,
                first_10_from_page[i].message_id,
                "Pagination should maintain same order as get_messages()"
            );
        }

        // Clean up
        tokio::fs::remove_dir_all(temp_dir).await.ok();
        Ok(())
    }

    #[tokio::test]
    async fn test_count_messages_nonexistent_session() -> Result<()> {
        let temp_dir = std::env::temp_dir().join(format!("archivist_test_{}", Uuid::now_v7()));
        let archivist = mk_archivist(temp_dir.clone()).await?;

        // Count messages for non-existent session (should return 0, not error)
        let nonexistent_scroll_id = Uuid::now_v7();
        let count = archivist.count_messages(nonexistent_scroll_id, None).await?;
        assert_eq!(count, 0, "Non-existent session should have 0 messages");

        // Clean up
        tokio::fs::remove_dir_all(temp_dir).await.ok();
        Ok(())
    }
}