poimen-memory/tests/it_chunking.rs

use mem_chunk::{chunks, ChunkPolicy, Boundary, FlushTrigger};
use mem_chunk::record_source::VecSource;
use mem_core::{Record, Provenance, Role};
use futures::stream::StreamExt;
use time::macros::datetime;

#[tokio::test]
async fn a1_no_record_is_split() {
    let records = vec![
        Record {
            role: Role::User,
            text: "First record".to_string(),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 0,
            },
        },
        Record {
            role: Role::Assistant,
            text: "Second record".to_string(),
            timestamp: datetime!(2024-08-20 12:00:01 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 1,
            },
        },
        Record {
            role: Role::User,
            text: "Third record".to_string(),
            timestamp: datetime!(2024-08-20 12:00:02 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 2,
            },
        },
    ];

    let original_records = records.clone();
    let source = VecSource(records);
    let policy = ChunkPolicy::default();
    let mut chunk_stream = chunks(source, policy);

    let mut all_chunk_records = Vec::new();
    while let Some(result) = chunk_stream.next().await {
        let chunk = result.unwrap();
        for record in &chunk.records {
            // Verify that this record appears in the original set
            assert!(original_records.iter().any(|r| {
                r.role == record.role && r.text == record.text
            }));
        }
        all_chunk_records.extend(chunk.records);
    }

    // Verify no record was split - all records should be intact
    for i in 0..original_records.len() {
        assert_eq!(all_chunk_records[i].text, original_records[i].text);
    }
}

#[tokio::test]
async fn a2_lossless() {
    let records = vec![
        Record {
            role: Role::User,
            text: "Hello".to_string(),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 0,
            },
        },
        Record {
            role: Role::Assistant,
            text: "Hi".to_string(),
            timestamp: datetime!(2024-08-20 12:00:01 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 1,
            },
        },
        Record {
            role: Role::User,
            text: "How are you?".to_string(),
            timestamp: datetime!(2024-08-20 12:00:02 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: 2,
            },
        },
    ];

    let original_count = records.len();
    let source = VecSource(records.clone());
    let policy = ChunkPolicy::default();
    let mut chunk_stream = chunks(source, policy);

    let mut all_chunk_records = Vec::new();
    while let Some(result) = chunk_stream.next().await {
        let chunk = result.unwrap();
        all_chunk_records.extend(chunk.records);
    }

    // Flatten all chunk records and assert sequence equals input
    assert_eq!(all_chunk_records.len(), original_count);
    for i in 0..original_count {
        assert_eq!(all_chunk_records[i].role, records[i].role);
        assert_eq!(all_chunk_records[i].text, records[i].text);
    }
}

#[tokio::test]
async fn a3_t_is_contiguous() {
    let records: Vec<Record> = (0..10)
        .map(|i| Record {
            role: Role::User,
            text: format!("Message {}", i),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: i as u64,
            },
        })
        .collect();

    let source = VecSource(records);
    let policy = ChunkPolicy {
        max_tokens: 10,  // Small budget to force multiple chunks
        split_on: Boundary::Record,
        flush: FlushTrigger::Tokens(10),
    };
    let mut chunk_stream = chunks(source, policy);

    let mut t_values = Vec::new();
    while let Some(result) = chunk_stream.next().await {
        let chunk = result.unwrap();
        t_values.push(chunk.t);
    }

    // Assert t values are exactly 1..=n
    assert!(!t_values.is_empty());
    for i in 0..t_values.len() {
        assert_eq!(t_values[i], (i + 1) as u32);
    }
}

#[tokio::test]
async fn a4_respects_budget() {
    let records: Vec<Record> = (0..5)
        .map(|i| Record {
            role: Role::User,
            text: "x".repeat(100).to_string(),  // ~25 tokens each
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "session1".to_string(),
                offset: i as u64,
            },
        })
        .collect();

    let source = VecSource(records);
    let budget = 75;  // 3 records worth
    let policy = ChunkPolicy {
        max_tokens: budget,
        split_on: Boundary::Record,
        flush: FlushTrigger::Tokens(budget),
    };
    let mut chunk_stream = chunks(source, policy);

    while let Some(result) = chunk_stream.next().await {
        let chunk = result.unwrap();
        // Each chunk should be under budget or have exactly one record
        if chunk.records.len() == 1 {
            // Single oversized record
        } else {
            // Multiple records should be under budget
            assert!(chunk.tokens <= budget);
        }
    }
}

#[tokio::test]
async fn a5_oversized_record_survives() {
    let oversized_record = Record {
        role: Role::ToolResult,
        text: "x".repeat(3000).to_string(),  // ~750 tokens, 10× the budget
        timestamp: datetime!(2024-08-20 12:00:00 UTC),
        provenance: Provenance {
            source_id: "session1".to_string(),
            offset: 0,
        },
    };

    let records = vec![oversized_record.clone()];
    let source = VecSource(records);
    let policy = ChunkPolicy {
        max_tokens: 100,  // Very small budget
        split_on: Boundary::Record,
        flush: FlushTrigger::Tokens(100),
    };
    let mut chunk_stream = chunks(source, policy);

    let chunk = chunk_stream.next().await.unwrap().unwrap();
    assert_eq!(chunk.records.len(), 1);
    assert_eq!(chunk.records[0].text, oversized_record.text);
    assert_eq!(chunk.records[0].role, oversized_record.role);
    
    // Verify the oversized record is not truncated
    assert!(chunk.records[0].text.len() >= 3000);
}

#[tokio::test]
async fn a6_empty_source() {
    let source = VecSource(vec![]);
    let policy = ChunkPolicy::default();
    let mut chunk_stream = chunks(source, policy);

    let result = chunk_stream.next().await;
    assert!(result.is_none(), "Empty source should yield no chunks");
}