use claudius::{MessageParam, MessageRole};
use crate::{AgentID, CHUNK_SIZE_LIMIT, MountID};
#[derive(Debug)]
pub struct TransactionSerializationError {
pub agent_id: AgentID,
pub context_seq_no: u32,
pub transaction_seq_no: u64,
pub operation: String,
pub error: serde_json::Error,
}
impl std::fmt::Display for TransactionSerializationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"Serialization error in {} for transaction {}/{}/{}: {}",
self.operation, self.agent_id, self.context_seq_no, self.transaction_seq_no, self.error
)
}
}
impl std::error::Error for TransactionSerializationError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(&self.error)
}
}
#[derive(Debug)]
pub enum FromChunksError {
NoChunks,
MismatchedTransaction {
agent_id: AgentID,
context_seq_no: u32,
transaction_seq_no: u64,
},
MissingChunks {
agent_id: AgentID,
context_seq_no: u32,
transaction_seq_no: u64,
expected: u32,
actual: u32,
},
ExtraChunks {
agent_id: AgentID,
context_seq_no: u32,
transaction_seq_no: u64,
expected: u32,
actual: u32,
},
InvalidSequence {
agent_id: AgentID,
context_seq_no: u32,
transaction_seq_no: u64,
expected: u32,
actual: u32,
},
Deserialization {
agent_id: AgentID,
context_seq_no: u32,
transaction_seq_no: u64,
error: serde_json::Error,
},
}
impl std::fmt::Display for FromChunksError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FromChunksError::NoChunks => write!(f, "No chunks provided to reconstruct transaction"),
FromChunksError::MismatchedTransaction {
agent_id,
context_seq_no,
transaction_seq_no,
} => {
write!(
f,
"Chunks belong to different transactions than {}/{}/{}",
agent_id, context_seq_no, transaction_seq_no
)
}
FromChunksError::MissingChunks {
agent_id,
context_seq_no,
transaction_seq_no,
expected,
actual,
} => {
write!(
f,
"Missing chunks for transaction {}/{}/{}: expected {}, got {}",
agent_id, context_seq_no, transaction_seq_no, expected, actual
)
}
FromChunksError::ExtraChunks {
agent_id,
context_seq_no,
transaction_seq_no,
expected,
actual,
} => {
write!(
f,
"Extra chunks for transaction {}/{}/{}: expected {}, got {}",
agent_id, context_seq_no, transaction_seq_no, expected, actual
)
}
FromChunksError::InvalidSequence {
agent_id,
context_seq_no,
transaction_seq_no,
expected,
actual,
} => {
write!(
f,
"Invalid chunk sequence for transaction {}/{}/{}: expected {}, got {}",
agent_id, context_seq_no, transaction_seq_no, expected, actual
)
}
FromChunksError::Deserialization {
agent_id,
context_seq_no,
transaction_seq_no,
error,
} => {
write!(
f,
"Failed to deserialize transaction {}/{}/{}: {}",
agent_id, context_seq_no, transaction_seq_no, error
)
}
}
}
}
impl std::error::Error for FromChunksError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
FromChunksError::Deserialization { error, .. } => Some(error),
_ => None,
}
}
}
#[derive(Debug)]
pub struct ChunkSizeExceededError {
pub item_type: String,
pub actual_size: usize,
pub limit: usize,
}
impl std::fmt::Display for ChunkSizeExceededError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{} size {} exceeds limit {}",
self.item_type, self.actual_size, self.limit
)
}
}
impl std::error::Error for ChunkSizeExceededError {}
#[derive(Debug)]
pub enum InvariantViolation {
ChunkSizeExceeded(ChunkSizeExceededError),
ConsecutiveMessagesWithSameRole {
agent_id: AgentID,
context_seq_no: u32,
transaction_seq_no: u64,
role: MessageRole,
positions: (usize, usize),
},
}
impl std::fmt::Display for InvariantViolation {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
InvariantViolation::ChunkSizeExceeded(err) => write!(f, "Chunk size exceeded: {}", err),
InvariantViolation::ConsecutiveMessagesWithSameRole {
agent_id,
context_seq_no,
transaction_seq_no,
role,
positions,
} => {
write!(
f,
"Consecutive messages with same role {:?} at positions {} and {} in transaction {}/{}/{}",
role, positions.0, positions.1, agent_id, context_seq_no, transaction_seq_no
)
}
}
}
}
impl std::error::Error for InvariantViolation {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
InvariantViolation::ChunkSizeExceeded(err) => Some(err),
_ => None,
}
}
}
#[derive(Clone, Debug, Default, serde::Deserialize, serde::Serialize)]
pub struct Transaction {
pub agent_id: AgentID,
pub context_seq_no: u32,
pub transaction_seq_no: u64,
pub msgs: Vec<MessageParam>,
pub writes: Vec<FileWrite>,
}
impl Transaction {
pub fn messages(&self) -> impl DoubleEndedIterator<Item = MessageParam> + '_ {
self.msgs.iter().cloned()
}
pub fn chunk_transaction(
&self,
) -> Result<Vec<TransactionChunk>, TransactionSerializationError> {
let mut serialized =
serde_json::to_string(self).map_err(|e| TransactionSerializationError {
agent_id: self.agent_id,
context_seq_no: self.context_seq_no,
transaction_seq_no: self.transaction_seq_no,
operation: "serialize".to_string(),
error: e,
})?;
if serialized.len() <= CHUNK_SIZE_LIMIT {
return Ok(vec![TransactionChunk {
agent_id: self.agent_id,
context_seq_no: self.context_seq_no,
transaction_seq_no: self.transaction_seq_no,
chunk_seq_no: 0,
total_chunks: 1,
data: serialized,
}]);
}
let mut chunks = Vec::new();
let mut chunk_seq_no = 0;
while !serialized.is_empty() {
let chunk_end = if serialized.len() <= CHUNK_SIZE_LIMIT {
serialized.len()
} else {
let mut end = CHUNK_SIZE_LIMIT;
while end > 0 && !serialized.is_char_boundary(end) {
end -= 1;
}
if end == 0 {
serialized.chars().next().unwrap().len_utf8()
} else {
end
}
};
let chunk = serialized[..chunk_end].to_string();
serialized = serialized[chunk_end..].to_string();
chunks.push(TransactionChunk {
agent_id: self.agent_id,
context_seq_no: self.context_seq_no,
transaction_seq_no: self.transaction_seq_no,
chunk_seq_no,
total_chunks: 0,
data: chunk,
});
chunk_seq_no += 1;
}
let total_chunks = chunks.len() as u32;
for chunk in &mut chunks {
chunk.total_chunks = total_chunks;
}
Ok(chunks)
}
pub fn check_invariants(&self) -> Result<(), InvariantViolation> {
for w in self.writes.iter() {
if w.data.len() >= CHUNK_SIZE_LIMIT {
return Err(InvariantViolation::ChunkSizeExceeded(
ChunkSizeExceededError {
item_type: "File write".to_string(),
actual_size: w.data.len(),
limit: CHUNK_SIZE_LIMIT,
},
));
}
}
for (i, window) in self.msgs.windows(2).enumerate() {
if window[0].role == window[1].role {
return Err(InvariantViolation::ConsecutiveMessagesWithSameRole {
agent_id: self.agent_id,
context_seq_no: self.context_seq_no,
transaction_seq_no: self.transaction_seq_no,
role: window[0].role,
positions: (i, i + 1),
});
}
}
Ok(())
}
}
impl Transaction {
pub fn from_chunks(chunks: Vec<TransactionChunk>) -> Result<Transaction, FromChunksError> {
if chunks.is_empty() {
return Err(FromChunksError::NoChunks);
}
let first = chunks[0].clone();
for chunk in &chunks {
if chunk.agent_id != first.agent_id
|| chunk.context_seq_no != first.context_seq_no
|| chunk.transaction_seq_no != first.transaction_seq_no
|| chunk.total_chunks != first.total_chunks
{
return Err(FromChunksError::MismatchedTransaction {
agent_id: first.agent_id,
context_seq_no: first.context_seq_no,
transaction_seq_no: first.transaction_seq_no,
});
}
}
let mut sorted_chunks = chunks;
sorted_chunks.sort_by_key(|chunk| chunk.chunk_seq_no);
for (i, chunk) in sorted_chunks.iter().enumerate() {
if chunk.chunk_seq_no != i as u32 {
return Err(FromChunksError::InvalidSequence {
agent_id: first.agent_id,
context_seq_no: first.context_seq_no,
transaction_seq_no: first.transaction_seq_no,
expected: i as u32,
actual: chunk.chunk_seq_no,
});
}
}
if sorted_chunks.len() != first.total_chunks as usize {
if sorted_chunks.len() > first.total_chunks as usize {
return Err(FromChunksError::ExtraChunks {
agent_id: first.agent_id,
context_seq_no: first.context_seq_no,
transaction_seq_no: first.transaction_seq_no,
expected: first.total_chunks,
actual: sorted_chunks.len() as u32,
});
} else {
return Err(FromChunksError::MissingChunks {
agent_id: first.agent_id,
context_seq_no: first.context_seq_no,
transaction_seq_no: first.transaction_seq_no,
expected: first.total_chunks,
actual: sorted_chunks.len() as u32,
});
}
}
let mut reconstructed_data = String::new();
for chunk in sorted_chunks {
reconstructed_data.push_str(&chunk.data);
}
let transaction: Transaction = serde_json::from_str(&reconstructed_data).map_err(|e| {
FromChunksError::Deserialization {
agent_id: first.agent_id,
context_seq_no: first.context_seq_no,
transaction_seq_no: first.transaction_seq_no,
error: e,
}
})?;
if transaction.agent_id != first.agent_id
|| transaction.context_seq_no != first.context_seq_no
|| transaction.transaction_seq_no != first.transaction_seq_no
{
return Err(FromChunksError::MismatchedTransaction {
agent_id: transaction.agent_id,
context_seq_no: transaction.context_seq_no,
transaction_seq_no: transaction.transaction_seq_no,
});
}
Ok(transaction)
}
}
#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
pub struct TransactionChunk {
pub agent_id: AgentID,
pub context_seq_no: u32,
pub transaction_seq_no: u64,
pub chunk_seq_no: u32,
pub total_chunks: u32,
pub data: String,
}
#[derive(Clone, Debug, Default, serde::Deserialize, serde::Serialize)]
pub struct FileWrite {
pub mount: MountID,
pub path: String,
pub data: String,
}
#[cfg(test)]
mod tests {
use super::*;
use claudius::MessageParamContent;
fn test_agent_id() -> AgentID {
AgentID::from_human_readable("agent:00000000-0000-0000-0000-000000000001").unwrap()
}
fn test_mount_id() -> MountID {
MountID::from_human_readable("mount:00000000-0000-0000-0000-000000000001").unwrap()
}
fn create_test_transaction() -> Transaction {
Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: "Hello".into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "Hi there!".into(),
},
],
writes: vec![FileWrite {
mount: test_mount_id(),
path: "test.txt".to_string(),
data: "test data".to_string(),
}],
}
}
#[test]
fn transaction_chunks_correctly_for_small_data() {
let transaction = create_test_transaction();
let chunks = transaction.chunk_transaction().unwrap();
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].chunk_seq_no, 0);
assert_eq!(chunks[0].total_chunks, 1);
assert_eq!(chunks[0].agent_id, transaction.agent_id);
assert_eq!(chunks[0].context_seq_no, transaction.context_seq_no);
assert_eq!(chunks[0].transaction_seq_no, transaction.transaction_seq_no);
}
#[test]
fn transaction_chunks_correctly_for_large_data() {
let mut transaction = create_test_transaction();
let large_content = "x".repeat(CHUNK_SIZE_LIMIT * 2);
transaction.msgs.push(MessageParam::new(
MessageParamContent::String(large_content),
MessageRole::User,
));
let chunks = transaction.chunk_transaction().unwrap();
assert!(chunks.len() > 1);
for (i, chunk) in chunks.iter().enumerate() {
assert_eq!(chunk.chunk_seq_no, i as u32);
assert_eq!(chunk.total_chunks, chunks.len() as u32);
assert_eq!(chunk.agent_id, transaction.agent_id);
}
}
#[test]
fn transaction_reconstructs_correctly_from_chunks() {
let original = create_test_transaction();
let chunks = original.chunk_transaction().unwrap();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(original.agent_id, reconstructed.agent_id);
assert_eq!(original.context_seq_no, reconstructed.context_seq_no);
assert_eq!(
original.transaction_seq_no,
reconstructed.transaction_seq_no
);
assert_eq!(original.msgs.len(), reconstructed.msgs.len());
assert_eq!(original.writes.len(), reconstructed.writes.len());
}
#[test]
fn from_chunks_fails_with_no_chunks() {
let result = Transaction::from_chunks(vec![]);
assert!(matches!(result, Err(FromChunksError::NoChunks)));
}
#[test]
fn from_chunks_fails_with_mismatched_transaction_ids() {
let chunk1 = TransactionChunk {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 1,
chunk_seq_no: 0,
total_chunks: 2,
data: "{}".to_string(),
};
let mut chunk2 = chunk1.clone();
chunk2.transaction_seq_no = 2; chunk2.chunk_seq_no = 1;
let result = Transaction::from_chunks(vec![chunk1, chunk2]);
assert!(matches!(
result,
Err(FromChunksError::MismatchedTransaction { .. })
));
}
#[test]
fn from_chunks_fails_with_missing_chunks() {
let chunk = TransactionChunk {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 1,
chunk_seq_no: 0,
total_chunks: 2, data: "{}".to_string(),
};
let result = Transaction::from_chunks(vec![chunk]);
assert!(matches!(result, Err(FromChunksError::MissingChunks { .. })));
}
#[test]
fn from_chunks_fails_with_extra_chunks() {
let chunks = (0..=6)
.map(|i| TransactionChunk {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
chunk_seq_no: i,
total_chunks: 6, data: format!("{{\"chunk\":{}}}", i),
})
.collect();
match Transaction::from_chunks(chunks) {
Err(FromChunksError::ExtraChunks {
expected, actual, ..
}) => {
assert_eq!(expected, 6);
assert_eq!(actual, 7);
}
_ => panic!("Expected ExtraChunks error"),
}
}
#[test]
fn check_invariants_passes_for_valid_transaction() {
let transaction = create_test_transaction();
assert!(transaction.check_invariants().is_ok());
}
#[test]
fn check_invariants_fails_for_consecutive_same_role_messages() {
let mut transaction = create_test_transaction();
transaction
.msgs
.push(MessageParam::assistant("Another assistant message"));
let result = transaction.check_invariants();
assert!(matches!(
result,
Err(InvariantViolation::ConsecutiveMessagesWithSameRole { .. })
));
if let Err(InvariantViolation::ConsecutiveMessagesWithSameRole {
role, positions, ..
}) = result
{
assert_eq!(role, MessageRole::Assistant); assert_eq!(positions, (1, 2)); }
}
#[test]
fn check_invariants_fails_for_oversized_file_write() {
let mut transaction = create_test_transaction();
transaction.writes[0].data = "x".repeat(CHUNK_SIZE_LIMIT + 1);
let result = transaction.check_invariants();
assert!(matches!(
result,
Err(InvariantViolation::ChunkSizeExceeded(_))
));
}
#[test]
fn messages_iterator_works() {
let transaction = create_test_transaction();
let messages: Vec<_> = transaction.messages().collect();
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].role, MessageRole::User);
assert_eq!(messages[1].role, MessageRole::Assistant);
}
#[test]
fn utf8_chunking_preserves_boundaries() {
let mut transaction = create_test_transaction();
let mut content = "a".repeat(CHUNK_SIZE_LIMIT - 10);
content.push_str("π¦π¦π¦π¦π¦"); content.push_str(&"b".repeat(100));
transaction.msgs.push(MessageParam::user(content));
let chunks = transaction.chunk_transaction().unwrap();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(transaction.msgs.len(), reconstructed.msgs.len());
for (orig, recon) in transaction.msgs.iter().zip(reconstructed.msgs.iter()) {
assert_eq!(orig.role, recon.role);
assert_eq!(orig.content, recon.content);
}
}
#[test]
fn check_invariants_valid_transaction() {
let transaction = create_test_transaction();
assert!(transaction.check_invariants().is_ok());
}
#[test]
fn check_invariants_consecutive_messages_same_role_user_user() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: "First user message".into(),
},
MessageParam {
role: MessageRole::User,
content: "Second user message".into(),
},
],
writes: vec![],
};
match transaction.check_invariants() {
Err(InvariantViolation::ConsecutiveMessagesWithSameRole {
agent_id,
context_seq_no,
transaction_seq_no,
role,
positions,
}) => {
assert_eq!(agent_id, test_agent_id());
assert_eq!(context_seq_no, 1);
assert_eq!(transaction_seq_no, 42);
assert_eq!(role, MessageRole::User);
assert_eq!(positions, (0, 1));
}
_ => panic!("Expected ConsecutiveMessagesWithSameRole error"),
}
}
#[test]
fn check_invariants_consecutive_messages_same_role_assistant_assistant() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 2,
transaction_seq_no: 100,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: "User message".into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "First assistant message".into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "Second assistant message".into(),
},
],
writes: vec![],
};
match transaction.check_invariants() {
Err(InvariantViolation::ConsecutiveMessagesWithSameRole {
agent_id,
context_seq_no,
transaction_seq_no,
role,
positions,
}) => {
assert_eq!(agent_id, test_agent_id());
assert_eq!(context_seq_no, 2);
assert_eq!(transaction_seq_no, 100);
assert_eq!(role, MessageRole::Assistant);
assert_eq!(positions, (1, 2));
}
_ => panic!("Expected ConsecutiveMessagesWithSameRole error"),
}
}
#[test]
fn check_invariants_multiple_consecutive_violations_reports_first() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: "First user message".into(),
},
MessageParam {
role: MessageRole::User,
content: "Second user message".into(),
},
MessageParam {
role: MessageRole::User,
content: "Third user message".into(),
},
],
writes: vec![],
};
match transaction.check_invariants() {
Err(InvariantViolation::ConsecutiveMessagesWithSameRole { positions, .. }) => {
assert_eq!(positions, (0, 1));
}
_ => panic!("Expected ConsecutiveMessagesWithSameRole error"),
}
}
#[test]
fn check_invariants_alternating_roles_valid() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: "User 1".into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "Assistant 1".into(),
},
MessageParam {
role: MessageRole::User,
content: "User 2".into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "Assistant 2".into(),
},
],
writes: vec![],
};
assert!(transaction.check_invariants().is_ok());
}
#[test]
fn check_invariants_single_message_valid() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![MessageParam {
role: MessageRole::User,
content: "Only message".into(),
}],
writes: vec![],
};
assert!(transaction.check_invariants().is_ok());
}
#[test]
fn check_invariants_empty_messages_valid() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![],
writes: vec![],
};
assert!(transaction.check_invariants().is_ok());
}
#[test]
fn check_invariants_file_write_exceeds_chunk_size() {
let large_data = "x".repeat(CHUNK_SIZE_LIMIT);
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![],
writes: vec![FileWrite {
mount: test_mount_id(),
path: "large.txt".to_string(),
data: large_data.clone(),
}],
};
match transaction.check_invariants() {
Err(InvariantViolation::ChunkSizeExceeded(error)) => {
assert_eq!(error.item_type, "File write");
assert_eq!(error.actual_size, large_data.len());
assert_eq!(error.limit, CHUNK_SIZE_LIMIT);
}
_ => panic!("Expected ChunkSizeExceeded error"),
}
}
#[test]
fn check_invariants_file_write_at_chunk_size_limit_valid() {
let data_at_limit = "x".repeat(CHUNK_SIZE_LIMIT - 1);
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![],
writes: vec![FileWrite {
mount: test_mount_id(),
path: "at_limit.txt".to_string(),
data: data_at_limit,
}],
};
assert!(transaction.check_invariants().is_ok());
}
#[test]
fn check_invariants_multiple_file_writes_one_exceeds() {
let large_data = "x".repeat(CHUNK_SIZE_LIMIT);
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![],
writes: vec![
FileWrite {
mount: test_mount_id(),
path: "small.txt".to_string(),
data: "small".to_string(),
},
FileWrite {
mount: test_mount_id(),
path: "large.txt".to_string(),
data: large_data,
},
],
};
match transaction.check_invariants() {
Err(InvariantViolation::ChunkSizeExceeded(_)) => {
}
_ => panic!("Expected ChunkSizeExceeded error"),
}
}
#[test]
fn chunk_transaction_small_transaction_single_chunk() {
let transaction = create_test_transaction();
let chunks = transaction.chunk_transaction().unwrap();
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].agent_id, transaction.agent_id);
assert_eq!(chunks[0].context_seq_no, transaction.context_seq_no);
assert_eq!(chunks[0].transaction_seq_no, transaction.transaction_seq_no);
assert_eq!(chunks[0].chunk_seq_no, 0);
assert_eq!(chunks[0].total_chunks, 1);
let deserialized: Transaction = serde_json::from_str(&chunks[0].data).unwrap();
assert_eq!(deserialized.agent_id, transaction.agent_id);
assert_eq!(deserialized.msgs.len(), transaction.msgs.len());
}
#[test]
fn chunk_transaction_large_transaction_multiple_chunks() {
let large_content = "x".repeat(CHUNK_SIZE_LIMIT);
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: large_content.into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "Response".into(),
},
],
writes: vec![],
};
let chunks = transaction.chunk_transaction().unwrap();
assert!(chunks.len() > 1);
for (i, chunk) in chunks.iter().enumerate() {
assert_eq!(chunk.agent_id, transaction.agent_id);
assert_eq!(chunk.context_seq_no, transaction.context_seq_no);
assert_eq!(chunk.transaction_seq_no, transaction.transaction_seq_no);
assert_eq!(chunk.chunk_seq_no, i as u32);
assert_eq!(chunk.total_chunks, chunks.len() as u32);
assert!(chunk.data.len() <= CHUNK_SIZE_LIMIT);
}
let reconstructed_data: String = chunks.iter().map(|c| c.data.as_str()).collect();
let deserialized: Transaction = serde_json::from_str(&reconstructed_data).unwrap();
assert_eq!(deserialized.agent_id, transaction.agent_id);
assert_eq!(deserialized.msgs.len(), transaction.msgs.len());
}
#[test]
fn chunk_transaction_exactly_at_limit() {
let mut transaction = create_test_transaction();
let serialized_base = serde_json::to_string(&transaction).unwrap();
let needed_padding = CHUNK_SIZE_LIMIT - serialized_base.len();
if needed_padding > 0 {
transaction.msgs[0].content = format!("Hello{}", "x".repeat(needed_padding - 5)).into();
}
let chunks = transaction.chunk_transaction().unwrap();
if serde_json::to_string(&transaction).unwrap().len() <= CHUNK_SIZE_LIMIT {
assert_eq!(chunks.len(), 1);
} else {
assert!(chunks.len() > 1);
}
}
#[test]
fn from_chunks_single_chunk() {
let original = create_test_transaction();
let chunks = original.chunk_transaction().unwrap();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(reconstructed.agent_id, original.agent_id);
assert_eq!(reconstructed.context_seq_no, original.context_seq_no);
assert_eq!(
reconstructed.transaction_seq_no,
original.transaction_seq_no
);
assert_eq!(reconstructed.msgs.len(), original.msgs.len());
assert_eq!(reconstructed.writes.len(), original.writes.len());
}
#[test]
fn from_chunks_multiple_chunks() {
let large_content = "x".repeat(CHUNK_SIZE_LIMIT);
let original = Transaction {
agent_id: test_agent_id(),
context_seq_no: 5,
transaction_seq_no: 123,
msgs: vec![MessageParam {
role: MessageRole::User,
content: large_content.into(),
}],
writes: vec![],
};
let chunks = original.chunk_transaction().unwrap();
assert!(chunks.len() > 1);
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(reconstructed.agent_id, original.agent_id);
assert_eq!(reconstructed.context_seq_no, original.context_seq_no);
assert_eq!(
reconstructed.transaction_seq_no,
original.transaction_seq_no
);
assert_eq!(reconstructed.msgs.len(), original.msgs.len());
}
#[test]
fn from_chunks_unordered_chunks() {
let large_content = "x".repeat(CHUNK_SIZE_LIMIT);
let original = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![MessageParam {
role: MessageRole::User,
content: large_content.into(),
}],
writes: vec![],
};
let mut chunks = original.chunk_transaction().unwrap();
assert!(chunks.len() > 1);
chunks.reverse();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(reconstructed.agent_id, original.agent_id);
}
#[test]
fn from_chunks_empty_chunks() {
match Transaction::from_chunks(vec![]) {
Err(FromChunksError::NoChunks) => {
}
_ => panic!("Expected NoChunks error"),
}
}
#[test]
fn from_chunks_mismatched_agent_id() {
let original = create_test_transaction();
let mut chunks = original.chunk_transaction().unwrap();
let different_agent_id =
AgentID::from_human_readable("agent:00000000-0000-0000-0000-000000000002").unwrap();
chunks[0].agent_id = different_agent_id;
match Transaction::from_chunks(chunks) {
Err(FromChunksError::MismatchedTransaction {
agent_id,
context_seq_no,
transaction_seq_no,
}) => {
assert_eq!(agent_id, original.agent_id);
assert_eq!(context_seq_no, original.context_seq_no);
assert_eq!(transaction_seq_no, original.transaction_seq_no);
}
_ => panic!("Expected MismatchedTransaction error"),
}
}
#[test]
fn from_chunks_mismatched_context_seq_no() {
let original = create_test_transaction();
let mut chunks = original.chunk_transaction().unwrap();
chunks[0].context_seq_no = 999;
match Transaction::from_chunks(chunks) {
Err(FromChunksError::MismatchedTransaction { .. }) => {
}
_ => panic!("Expected MismatchedTransaction error"),
}
}
#[test]
fn from_chunks_mismatched_transaction_seq_no() {
let original = create_test_transaction();
let mut chunks = original.chunk_transaction().unwrap();
chunks[0].transaction_seq_no = 999;
match Transaction::from_chunks(chunks) {
Err(FromChunksError::MismatchedTransaction { .. }) => {
}
_ => panic!("Expected MismatchedTransaction error"),
}
}
#[test]
fn from_chunks_mismatched_total_chunks() {
let original = create_test_transaction();
let mut chunks = original.chunk_transaction().unwrap();
chunks[0].total_chunks = 999;
match Transaction::from_chunks(chunks) {
Err(FromChunksError::MissingChunks {
expected, actual, ..
}) => {
assert_eq!(expected, 999);
assert_eq!(actual, 1);
}
_ => panic!("Expected MissingChunks error"),
}
}
#[test]
fn from_chunks_missing_chunks() {
let large_content = "x".repeat(CHUNK_SIZE_LIMIT);
let original = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![MessageParam {
role: MessageRole::User,
content: large_content.into(),
}],
writes: vec![],
};
let mut chunks = original.chunk_transaction().unwrap();
assert!(chunks.len() > 1);
let expected_count = chunks.len();
chunks.pop();
match Transaction::from_chunks(chunks) {
Err(FromChunksError::MissingChunks {
expected, actual, ..
}) => {
assert_eq!(expected, expected_count as u32);
assert_eq!(actual, (expected_count - 1) as u32);
}
_ => panic!("Expected MissingChunks error"),
}
}
#[test]
fn from_chunks_invalid_sequence_gap() {
let large_content = "x".repeat(CHUNK_SIZE_LIMIT * 3);
let original = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![MessageParam {
role: MessageRole::User,
content: large_content.into(),
}],
writes: vec![],
};
let mut chunks = original.chunk_transaction().unwrap();
assert!(chunks.len() > 2);
chunks[1].chunk_seq_no = 5;
match Transaction::from_chunks(chunks) {
Err(FromChunksError::InvalidSequence {
expected, actual, ..
}) => {
assert_eq!(expected, 1);
assert_eq!(actual, 2);
}
_ => panic!("Expected InvalidSequence error"),
}
}
#[test]
fn from_chunks_invalid_sequence_duplicate() {
let large_content = "x".repeat(CHUNK_SIZE_LIMIT);
let original = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![MessageParam {
role: MessageRole::User,
content: large_content.into(),
}],
writes: vec![],
};
let mut chunks = original.chunk_transaction().unwrap();
assert!(chunks.len() > 1);
chunks[1].chunk_seq_no = 0;
match Transaction::from_chunks(chunks) {
Err(FromChunksError::InvalidSequence {
expected, actual, ..
}) => {
assert_eq!(expected, 1);
assert_eq!(actual, 0);
}
_ => panic!("Expected InvalidSequence error"),
}
}
#[test]
fn from_chunks_corrupted_json_data() {
let original = create_test_transaction();
let mut chunks = original.chunk_transaction().unwrap();
chunks[0].data = "invalid json {".to_string();
match Transaction::from_chunks(chunks) {
Err(FromChunksError::Deserialization { .. }) => {
}
_ => panic!("Expected Deserialization error"),
}
}
#[test]
fn messages_iterator_forward() {
let transaction = create_test_transaction();
let messages: Vec<MessageParam> = transaction.messages().collect();
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].role, MessageRole::User);
assert_eq!(messages[1].role, MessageRole::Assistant);
}
#[test]
fn messages_iterator_reverse() {
let transaction = create_test_transaction();
let messages: Vec<MessageParam> = transaction.messages().rev().collect();
assert_eq!(messages.len(), 2);
assert_eq!(messages[0].role, MessageRole::Assistant);
assert_eq!(messages[1].role, MessageRole::User);
}
#[test]
fn messages_iterator_empty() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![],
writes: vec![],
};
let messages: Vec<MessageParam> = transaction.messages().collect();
assert_eq!(messages.len(), 0);
}
#[test]
fn transaction_with_unicode_content() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 1,
transaction_seq_no: 42,
msgs: vec![
MessageParam {
role: MessageRole::User,
content: "Hello π δΈη π".into(),
},
MessageParam {
role: MessageRole::Assistant,
content: "γγγ«γ‘γ―! π".into(),
},
],
writes: vec![],
};
assert!(transaction.check_invariants().is_ok());
let chunks = transaction.chunk_transaction().unwrap();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
match &reconstructed.msgs[0].content {
claudius::MessageParamContent::String(s) => assert_eq!(s, "Hello π δΈη π"),
_ => panic!("Expected String content"),
}
match &reconstructed.msgs[1].content {
claudius::MessageParamContent::String(s) => assert_eq!(s, "γγγ«γ‘γ―! π"),
_ => panic!("Expected String content"),
}
}
#[test]
fn transaction_zero_sequence_numbers() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: 0,
transaction_seq_no: 0,
msgs: vec![],
writes: vec![],
};
assert!(transaction.check_invariants().is_ok());
let chunks = transaction.chunk_transaction().unwrap();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(reconstructed.context_seq_no, 0);
assert_eq!(reconstructed.transaction_seq_no, 0);
}
#[test]
fn transaction_max_sequence_numbers() {
let transaction = Transaction {
agent_id: test_agent_id(),
context_seq_no: u32::MAX,
transaction_seq_no: u64::MAX,
msgs: vec![],
writes: vec![],
};
assert!(transaction.check_invariants().is_ok());
let chunks = transaction.chunk_transaction().unwrap();
let reconstructed = Transaction::from_chunks(chunks).unwrap();
assert_eq!(reconstructed.context_seq_no, u32::MAX);
assert_eq!(reconstructed.transaction_seq_no, u64::MAX);
}
}