enya-agent 0.1.0

package/crates/enact-core/src/kernel/persistence/message_store.rs

@@ -0,0 +1,908 @@
+//! MessageStore - User-facing message persistence
+//!
+//! The MessageStore handles thread/message persistence for UI display.
+//! Unlike EventStore (audit-focused, immutable), MessageStore is:
+//! - **Mutable**: Messages can be soft-deleted/edited
+//! - **User-facing**: AI SDK compatible format
+//! - **Hierarchical**: Thread → Message with parent linkage
+//!
+//! ## Guarantees
+//!
+//! - **Soft Delete**: Messages are never hard-deleted (GDPR compliance)
+//! - **Thread Ownership**: Messages belong to exactly one thread
+//! - **Parent Linkage**: Every message tracks its causal origin
+//!
+//! @see docs/TECHNICAL/20-MESSAGE-PERSISTENCE-STRATEGY.md
+
+use async_trait::async_trait;
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::RwLock;
+
+use crate::kernel::{ExecutionId, MessageId, ParentType, TenantId, ThreadId, UserId};
+
+use super::StorageBackend;
+
+// =============================================================================
+// Thread - Conversation container
+// =============================================================================
+
+/// A conversation thread containing messages
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Thread {
+    /// Unique thread identifier
+    pub id: ThreadId,
+    /// Tenant for multi-tenancy
+    pub tenant_id: TenantId,
+    /// User who owns this thread
+    pub user_id: UserId,
+    /// Optional thread title (auto-generated from first message)
+    pub title: Option<String>,
+    /// When the thread was created
+    pub created_at: DateTime<Utc>,
+    /// When the thread was last updated
+    pub updated_at: DateTime<Utc>,
+    /// Soft delete timestamp (None = active)
+    pub deleted_at: Option<DateTime<Utc>>,
+}
+
+impl Thread {
+    /// Create a new thread
+    pub fn new(tenant_id: TenantId, user_id: UserId) -> Self {
+        let now = Utc::now();
+        Self {
+            id: ThreadId::new(),
+            tenant_id,
+            user_id,
+            title: None,
+            created_at: now,
+            updated_at: now,
+            deleted_at: None,
+        }
+    }
+
+    /// Create a thread with a specific ID (for testing/restore)
+    pub fn with_id(id: ThreadId, tenant_id: TenantId, user_id: UserId) -> Self {
+        let now = Utc::now();
+        Self {
+            id,
+            tenant_id,
+            user_id,
+            title: None,
+            created_at: now,
+            updated_at: now,
+            deleted_at: None,
+        }
+    }
+
+    /// Check if thread is soft-deleted
+    pub fn is_deleted(&self) -> bool {
+        self.deleted_at.is_some()
+    }
+}
+
+// =============================================================================
+// Message Role
+// =============================================================================
+
+/// Message role (AI SDK compatible)
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum MessageRole {
+    /// User message
+    User,
+    /// Assistant/AI response
+    Assistant,
+    /// System message (instructions, context)
+    System,
+}
+
+impl std::fmt::Display for MessageRole {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            MessageRole::User => write!(f, "user"),
+            MessageRole::Assistant => write!(f, "assistant"),
+            MessageRole::System => write!(f, "system"),
+        }
+    }
+}
+
+// =============================================================================
+// Message Parts (AI SDK compatible)
+// =============================================================================
+
+/// Message part for rich content (AI SDK compatible)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "kebab-case")]
+pub enum MessagePart {
+    /// Plain text content
+    Text { text: String },
+
+    /// Reasoning/thinking (collapsed in UI)
+    Reasoning { text: String },
+
+    /// Tool invocation
+    ToolCall {
+        tool_call_id: String,
+        tool_name: String,
+        args: serde_json::Value,
+    },
+
+    /// Tool result
+    ToolResult {
+        tool_call_id: String,
+        tool_name: String,
+        result: serde_json::Value,
+        is_error: bool,
+    },
+
+    /// Source/citation
+    Source {
+        source_id: String,
+        url: Option<String>,
+        title: Option<String>,
+    },
+
+    /// File attachment
+    File {
+        file_id: String,
+        filename: String,
+        mime_type: String,
+        size_bytes: u64,
+    },
+
+    /// Image (inline or referenced)
+    Image {
+        image_id: String,
+        url: Option<String>,
+        alt_text: Option<String>,
+    },
+
+    /// Code block
+    Code {
+        language: Option<String>,
+        code: String,
+    },
+}
+
+// =============================================================================
+// Token Usage
+// =============================================================================
+
+/// Token usage for billing and analytics
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct TokenUsage {
+    /// Input/prompt tokens
+    pub prompt_tokens: u32,
+    /// Output/completion tokens
+    pub completion_tokens: u32,
+    /// Total tokens (may include cached)
+    pub total_tokens: u32,
+}
+
+impl TokenUsage {
+    /// Create new token usage
+    pub fn new(prompt: u32, completion: u32) -> Self {
+        Self {
+            prompt_tokens: prompt,
+            completion_tokens: completion,
+            total_tokens: prompt + completion,
+        }
+    }
+}
+
+// =============================================================================
+// Execution Stats
+// =============================================================================
+
+/// Execution statistics for a message
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct ExecutionStats {
+    /// Total LLM invocations
+    pub llm_calls: u32,
+    /// Total tool invocations
+    pub tool_calls: u32,
+    /// Nested agent executions
+    pub sub_agents: u32,
+    /// Total steps executed
+    pub steps: u32,
+    /// Decisions made (audit)
+    pub decisions: u32,
+    /// Artifacts produced
+    pub artifacts: u32,
+}
+
+// =============================================================================
+// Cost Info
+// =============================================================================
+
+/// Cost information for billing
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CostInfo {
+    /// Input token cost
+    pub input_cost: f64,
+    /// Output token cost
+    pub output_cost: f64,
+    /// Total cost
+    pub total_cost: f64,
+    /// Currency (default: USD)
+    pub currency: String,
+}
+
+impl Default for CostInfo {
+    fn default() -> Self {
+        Self {
+            input_cost: 0.0,
+            output_cost: 0.0,
+            total_cost: 0.0,
+            currency: "USD".to_string(),
+        }
+    }
+}
+
+// =============================================================================
+// Finish Reason
+// =============================================================================
+
+/// Why the response finished
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum FinishReason {
+    /// Normal completion
+    Stop,
+    /// Hit max tokens
+    Length,
+    /// Needs to call tools
+    ToolCalls,
+    /// Content filter triggered
+    ContentFilter,
+    /// Error occurred
+    Error,
+}
+
+// =============================================================================
+// Message Metadata
+// =============================================================================
+
+/// Message metadata (AI SDK compatible)
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct MessageMetadata {
+    /// When response completed (Unix ms)
+    pub completed_at: Option<i64>,
+    /// Total duration in ms
+    pub duration_ms: Option<u64>,
+
+    /// Model used (e.g., "gpt-4o", "claude-3-opus")
+    pub model: Option<String>,
+    /// Provider (e.g., "azure", "anthropic")
+    pub provider: Option<String>,
+
+    /// Token usage (aggregated across all LLM calls)
+    pub token_usage: Option<TokenUsage>,
+
+    /// Execution statistics
+    pub stats: Option<ExecutionStats>,
+
+    /// Why the response finished
+    pub finish_reason: Option<FinishReason>,
+
+    /// Cost information (if available)
+    pub cost: Option<CostInfo>,
+}
+
+// =============================================================================
+// Message
+// =============================================================================
+
+/// A message in a thread (AI SDK compatible)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Message {
+    /// Unique message identifier
+    pub id: MessageId,
+    /// Thread this message belongs to
+    pub thread_id: ThreadId,
+    /// Link to execution (for assistant messages)
+    pub execution_id: Option<ExecutionId>,
+
+    // Hierarchy
+    /// Parent message (for reply chains)
+    pub parent_id: Option<MessageId>,
+    /// What triggered this message
+    pub parent_type: ParentType,
+
+    // Content
+    /// Message role
+    pub role: MessageRole,
+    /// Simple text content (for display)
+    pub content: String,
+    /// Rich content parts (AI SDK format)
+    pub parts: Vec<MessagePart>,
+
+    // Timestamps
+    /// When created
+    pub created_at: DateTime<Utc>,
+    /// When last updated
+    pub updated_at: Option<DateTime<Utc>>,
+    /// Soft delete timestamp
+    pub deleted_at: Option<DateTime<Utc>>,
+
+    /// Additional metadata
+    pub metadata: MessageMetadata,
+}
+
+impl Message {
+    /// Create a new user message
+    pub fn user(thread_id: ThreadId, content: impl Into<String>) -> Self {
+        let content = content.into();
+        Self {
+            id: MessageId::new(),
+            thread_id,
+            execution_id: None,
+            parent_id: None,
+            parent_type: ParentType::UserMessage,
+            role: MessageRole::User,
+            content: content.clone(),
+            parts: vec![MessagePart::Text { text: content }],
+            created_at: Utc::now(),
+            updated_at: None,
+            deleted_at: None,
+            metadata: MessageMetadata::default(),
+        }
+    }
+
+    /// Create a new assistant message
+    pub fn assistant(
+        thread_id: ThreadId,
+        execution_id: ExecutionId,
+        content: impl Into<String>,
+        parent_id: Option<MessageId>,
+    ) -> Self {
+        let content = content.into();
+        Self {
+            id: MessageId::new(),
+            thread_id,
+            execution_id: Some(execution_id),
+            parent_id,
+            parent_type: ParentType::UserMessage,
+            role: MessageRole::Assistant,
+            content: content.clone(),
+            parts: vec![MessagePart::Text { text: content }],
+            created_at: Utc::now(),
+            updated_at: None,
+            deleted_at: None,
+            metadata: MessageMetadata::default(),
+        }
+    }
+
+    /// Create a new system message
+    pub fn system(thread_id: ThreadId, content: impl Into<String>) -> Self {
+        let content = content.into();
+        Self {
+            id: MessageId::new(),
+            thread_id,
+            execution_id: None,
+            parent_id: None,
+            parent_type: ParentType::System,
+            role: MessageRole::System,
+            content: content.clone(),
+            parts: vec![MessagePart::Text { text: content }],
+            created_at: Utc::now(),
+            updated_at: None,
+            deleted_at: None,
+            metadata: MessageMetadata::default(),
+        }
+    }
+
+    /// Check if message is soft-deleted
+    pub fn is_deleted(&self) -> bool {
+        self.deleted_at.is_some()
+    }
+
+    /// Set the parent message
+    pub fn with_parent(mut self, parent_id: MessageId, parent_type: ParentType) -> Self {
+        self.parent_id = Some(parent_id);
+        self.parent_type = parent_type;
+        self
+    }
+
+    /// Add rich content parts
+    pub fn with_parts(mut self, parts: Vec<MessagePart>) -> Self {
+        self.parts = parts;
+        self
+    }
+
+    /// Set metadata
+    pub fn with_metadata(mut self, metadata: MessageMetadata) -> Self {
+        self.metadata = metadata;
+        self
+    }
+}
+
+// =============================================================================
+// MessageStore Trait
+// =============================================================================
+
+/// MessageStore trait - thread/message persistence
+///
+/// This is the core persistence trait for user-facing message storage.
+/// Unlike EventStore (immutable audit log), MessageStore supports:
+/// - Soft delete (GDPR compliance)
+/// - Message editing (optional)
+/// - Thread management
+#[async_trait]
+pub trait MessageStore: StorageBackend {
+    // =========================================================================
+    // Thread Operations
+    // =========================================================================
+
+    /// Create a new thread
+    async fn create_thread(&self, thread: Thread) -> anyhow::Result<ThreadId>;
+
+    /// Get a thread by ID
+    async fn get_thread(&self, thread_id: &ThreadId) -> anyhow::Result<Option<Thread>>;
+
+    /// Update thread (title, etc.)
+    async fn update_thread(&self, thread: Thread) -> anyhow::Result<()>;
+
+    /// Soft-delete a thread (sets deleted_at)
+    async fn delete_thread(&self, thread_id: &ThreadId) -> anyhow::Result<()>;
+
+    /// List threads for a user
+    async fn list_threads(
+        &self,
+        tenant_id: &TenantId,
+        user_id: &UserId,
+        limit: usize,
+        offset: usize,
+    ) -> anyhow::Result<Vec<Thread>>;
+
+    // =========================================================================
+    // Message Operations
+    // =========================================================================
+
+    /// Create a new message in a thread
+    async fn create_message(&self, message: Message) -> anyhow::Result<MessageId>;
+
+    /// Get a message by ID
+    async fn get_message(&self, message_id: &MessageId) -> anyhow::Result<Option<Message>>;
+
+    /// Update a message (content, metadata)
+    async fn update_message(&self, message: Message) -> anyhow::Result<()>;
+
+    /// Soft-delete a message
+    async fn delete_message(&self, message_id: &MessageId) -> anyhow::Result<()>;
+
+    /// Get all messages in a thread (ordered by created_at)
+    async fn list_messages(
+        &self,
+        thread_id: &ThreadId,
+        include_deleted: bool,
+    ) -> anyhow::Result<Vec<Message>>;
+
+    /// Get messages by execution ID
+    async fn get_messages_by_execution(
+        &self,
+        execution_id: &ExecutionId,
+    ) -> anyhow::Result<Vec<Message>>;
+
+    // =========================================================================
+    // Convenience Methods
+    // =========================================================================
+
+    /// Get or create a thread (implicit thread creation)
+    async fn get_or_create_thread(
+        &self,
+        tenant_id: TenantId,
+        user_id: UserId,
+    ) -> anyhow::Result<Thread> {
+        // Default: just create a new thread
+        let thread = Thread::new(tenant_id, user_id);
+        self.create_thread(thread.clone()).await?;
+        Ok(thread)
+    }
+
+    /// Count messages in a thread
+    async fn count_messages(&self, thread_id: &ThreadId) -> anyhow::Result<u64> {
+        let messages = self.list_messages(thread_id, false).await?;
+        Ok(messages.len() as u64)
+    }
+}
+
+// =============================================================================
+// InMemoryMessageStore
+// =============================================================================
+
+/// In-memory implementation of MessageStore
+///
+/// Thread-safe storage using RwLock. Suitable for single-node deployments
+/// and testing. For persistence across restarts, use SQLite implementation.
+#[derive(Default)]
+pub struct InMemoryMessageStore {
+    threads: RwLock<HashMap<String, Thread>>,
+    messages: RwLock<HashMap<String, Message>>,
+}
+
+impl InMemoryMessageStore {
+    /// Create a new empty message store
+    pub fn new() -> Self {
+        Self {
+            threads: RwLock::new(HashMap::new()),
+            messages: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Create an Arc-wrapped instance for sharing
+    pub fn shared() -> std::sync::Arc<Self> {
+        std::sync::Arc::new(Self::new())
+    }
+}
+
+#[async_trait]
+impl StorageBackend for InMemoryMessageStore {
+    fn name(&self) -> &str {
+        "in-memory-message-store"
+    }
+
+    fn requires_network(&self) -> bool {
+        false
+    }
+
+    async fn health_check(&self) -> anyhow::Result<()> {
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl MessageStore for InMemoryMessageStore {
+    async fn create_thread(&self, thread: Thread) -> anyhow::Result<ThreadId> {
+        let id = thread.id.clone();
+        let mut guard = self.threads.write().expect("lock poisoned");
+        guard.insert(id.to_string(), thread);
+        Ok(id)
+    }
+
+    async fn get_thread(&self, thread_id: &ThreadId) -> anyhow::Result<Option<Thread>> {
+        let guard = self.threads.read().expect("lock poisoned");
+        Ok(guard.get(&thread_id.to_string()).cloned())
+    }
+
+    async fn update_thread(&self, thread: Thread) -> anyhow::Result<()> {
+        let mut guard = self.threads.write().expect("lock poisoned");
+        guard.insert(thread.id.to_string(), thread);
+        Ok(())
+    }
+
+    async fn delete_thread(&self, thread_id: &ThreadId) -> anyhow::Result<()> {
+        let mut guard = self.threads.write().expect("lock poisoned");
+        if let Some(thread) = guard.get_mut(&thread_id.to_string()) {
+            thread.deleted_at = Some(Utc::now());
+        }
+        Ok(())
+    }
+
+    async fn list_threads(
+        &self,
+        tenant_id: &TenantId,
+        user_id: &UserId,
+        limit: usize,
+        offset: usize,
+    ) -> anyhow::Result<Vec<Thread>> {
+        let guard = self.threads.read().expect("lock poisoned");
+        let mut threads: Vec<_> = guard
+            .values()
+            .filter(|t| {
+                t.tenant_id == *tenant_id
+                    && t.user_id == *user_id
+                    && t.deleted_at.is_none()
+            })
+            .cloned()
+            .collect();
+
+        // Sort by created_at descending (newest first)
+        threads.sort_by(|a, b| b.created_at.cmp(&a.created_at));
+
+        Ok(threads.into_iter().skip(offset).take(limit).collect())
+    }
+
+    async fn create_message(&self, message: Message) -> anyhow::Result<MessageId> {
+        let id = message.id.clone();
+        let thread_id = message.thread_id.clone();
+
+        // Update thread's updated_at
+        {
+            let mut thread_guard = self.threads.write().expect("lock poisoned");
+            if let Some(thread) = thread_guard.get_mut(&thread_id.to_string()) {
+                thread.updated_at = Utc::now();
+            }
+        }
+
+        let mut guard = self.messages.write().expect("lock poisoned");
+        guard.insert(id.to_string(), message);
+        Ok(id)
+    }
+
+    async fn get_message(&self, message_id: &MessageId) -> anyhow::Result<Option<Message>> {
+        let guard = self.messages.read().expect("lock poisoned");
+        Ok(guard.get(&message_id.to_string()).cloned())
+    }
+
+    async fn update_message(&self, mut message: Message) -> anyhow::Result<()> {
+        message.updated_at = Some(Utc::now());
+        let mut guard = self.messages.write().expect("lock poisoned");
+        guard.insert(message.id.to_string(), message);
+        Ok(())
+    }
+
+    async fn delete_message(&self, message_id: &MessageId) -> anyhow::Result<()> {
+        let mut guard = self.messages.write().expect("lock poisoned");
+        if let Some(message) = guard.get_mut(&message_id.to_string()) {
+            message.deleted_at = Some(Utc::now());
+        }
+        Ok(())
+    }
+
+    async fn list_messages(
+        &self,
+        thread_id: &ThreadId,
+        include_deleted: bool,
+    ) -> anyhow::Result<Vec<Message>> {
+        let guard = self.messages.read().expect("lock poisoned");
+        let mut messages: Vec<_> = guard
+            .values()
+            .filter(|m| {
+                m.thread_id == *thread_id
+                    && (include_deleted || m.deleted_at.is_none())
+            })
+            .cloned()
+            .collect();
+
+        // Sort by created_at ascending
+        messages.sort_by(|a, b| a.created_at.cmp(&b.created_at));
+
+        Ok(messages)
+    }
+
+    async fn get_messages_by_execution(
+        &self,
+        execution_id: &ExecutionId,
+    ) -> anyhow::Result<Vec<Message>> {
+        let guard = self.messages.read().expect("lock poisoned");
+        let messages: Vec<_> = guard
+            .values()
+            .filter(|m| m.execution_id.as_ref() == Some(execution_id))
+            .cloned()
+            .collect();
+        Ok(messages)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_tenant() -> TenantId {
+        TenantId::from_string("tenant_test")
+    }
+
+    fn test_user() -> UserId {
+        UserId::from_string("user_test")
+    }
+
+    #[tokio::test]
+    async fn test_create_thread() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let id = thread.id.clone();
+
+        store.create_thread(thread).await.unwrap();
+
+        let loaded = store.get_thread(&id).await.unwrap();
+        assert!(loaded.is_some());
+        assert_eq!(loaded.unwrap().id, id);
+    }
+
+    #[tokio::test]
+    async fn test_soft_delete_thread() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let id = thread.id.clone();
+
+        store.create_thread(thread).await.unwrap();
+        store.delete_thread(&id).await.unwrap();
+
+        let loaded = store.get_thread(&id).await.unwrap().unwrap();
+        assert!(loaded.is_deleted());
+    }
+
+    #[tokio::test]
+    async fn test_list_threads_excludes_deleted() {
+        let store = InMemoryMessageStore::new();
+        let tenant = test_tenant();
+        let user = test_user();
+
+        // Create two threads
+        let thread1 = Thread::new(tenant.clone(), user.clone());
+        let thread2 = Thread::new(tenant.clone(), user.clone());
+        let id2 = thread2.id.clone();
+
+        store.create_thread(thread1).await.unwrap();
+        store.create_thread(thread2).await.unwrap();
+
+        // Delete one
+        store.delete_thread(&id2).await.unwrap();
+
+        // List should only return 1
+        let threads = store.list_threads(&tenant, &user, 100, 0).await.unwrap();
+        assert_eq!(threads.len(), 1);
+    }
+
+    #[tokio::test]
+    async fn test_create_message() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        store.create_thread(thread).await.unwrap();
+
+        let message = Message::user(thread_id.clone(), "Hello!");
+        let msg_id = message.id.clone();
+        store.create_message(message).await.unwrap();
+
+        let loaded = store.get_message(&msg_id).await.unwrap();
+        assert!(loaded.is_some());
+        assert_eq!(loaded.unwrap().content, "Hello!");
+    }
+
+    #[tokio::test]
+    async fn test_message_parent_chain() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        store.create_thread(thread).await.unwrap();
+
+        // User message
+        let user_msg = Message::user(thread_id.clone(), "What's the weather?");
+        let user_msg_id = user_msg.id.clone();
+        store.create_message(user_msg).await.unwrap();
+
+        // Assistant response
+        let exec_id = ExecutionId::new();
+        let assistant_msg = Message::assistant(
+            thread_id.clone(),
+            exec_id,
+            "The weather is sunny.",
+            Some(user_msg_id.clone()),
+        );
+        store.create_message(assistant_msg).await.unwrap();
+
+        // List messages
+        let messages = store.list_messages(&thread_id, false).await.unwrap();
+        assert_eq!(messages.len(), 2);
+        assert_eq!(messages[0].role, MessageRole::User);
+        assert_eq!(messages[1].role, MessageRole::Assistant);
+        assert_eq!(messages[1].parent_id, Some(user_msg_id));
+    }
+
+    #[tokio::test]
+    async fn test_soft_delete_message() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        store.create_thread(thread).await.unwrap();
+
+        let message = Message::user(thread_id.clone(), "Delete me");
+        let msg_id = message.id.clone();
+        store.create_message(message).await.unwrap();
+
+        store.delete_message(&msg_id).await.unwrap();
+
+        // Should not appear in normal list
+        let messages = store.list_messages(&thread_id, false).await.unwrap();
+        assert_eq!(messages.len(), 0);
+
+        // Should appear with include_deleted
+        let all_messages = store.list_messages(&thread_id, true).await.unwrap();
+        assert_eq!(all_messages.len(), 1);
+        assert!(all_messages[0].is_deleted());
+    }
+
+    #[tokio::test]
+    async fn test_get_messages_by_execution() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        store.create_thread(thread).await.unwrap();
+
+        let exec_id = ExecutionId::new();
+        let msg = Message::assistant(thread_id, exec_id.clone(), "Response", None);
+        store.create_message(msg).await.unwrap();
+
+        let messages = store.get_messages_by_execution(&exec_id).await.unwrap();
+        assert_eq!(messages.len(), 1);
+        assert_eq!(messages[0].execution_id, Some(exec_id));
+    }
+
+    #[tokio::test]
+    async fn test_message_with_parts() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        store.create_thread(thread).await.unwrap();
+
+        let exec_id = ExecutionId::new();
+        let parts = vec![
+            MessagePart::Reasoning {
+                text: "Let me think...".to_string(),
+            },
+            MessagePart::ToolCall {
+                tool_call_id: "tc_123".to_string(),
+                tool_name: "get_weather".to_string(),
+                args: serde_json::json!({"city": "NYC"}),
+            },
+            MessagePart::Text {
+                text: "The weather is sunny.".to_string(),
+            },
+        ];
+
+        let msg = Message::assistant(thread_id, exec_id, "The weather is sunny.", None)
+            .with_parts(parts);
+        let msg_id = msg.id.clone();
+        store.create_message(msg).await.unwrap();
+
+        let loaded = store.get_message(&msg_id).await.unwrap().unwrap();
+        assert_eq!(loaded.parts.len(), 3);
+    }
+
+    #[tokio::test]
+    async fn test_message_metadata() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        store.create_thread(thread).await.unwrap();
+
+        let exec_id = ExecutionId::new();
+        let metadata = MessageMetadata {
+            model: Some("gpt-4o".to_string()),
+            provider: Some("azure".to_string()),
+            duration_ms: Some(1500),
+            token_usage: Some(TokenUsage::new(100, 200)),
+            stats: Some(ExecutionStats {
+                llm_calls: 2,
+                tool_calls: 1,
+                sub_agents: 0,
+                steps: 3,
+                decisions: 1,
+                artifacts: 0,
+            }),
+            finish_reason: Some(FinishReason::Stop),
+            ..Default::default()
+        };
+
+        let msg = Message::assistant(thread_id, exec_id, "Response", None)
+            .with_metadata(metadata);
+        let msg_id = msg.id.clone();
+        store.create_message(msg).await.unwrap();
+
+        let loaded = store.get_message(&msg_id).await.unwrap().unwrap();
+        assert_eq!(loaded.metadata.model, Some("gpt-4o".to_string()));
+        assert_eq!(loaded.metadata.stats.as_ref().unwrap().llm_calls, 2);
+    }
+
+    #[tokio::test]
+    async fn test_thread_updated_on_message() {
+        let store = InMemoryMessageStore::new();
+        let thread = Thread::new(test_tenant(), test_user());
+        let thread_id = thread.id.clone();
+        let original_updated = thread.updated_at;
+        store.create_thread(thread).await.unwrap();
+
+        // Small delay to ensure timestamp differs
+        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+
+        let message = Message::user(thread_id.clone(), "Hello!");
+        store.create_message(message).await.unwrap();
+
+        let loaded_thread = store.get_thread(&thread_id).await.unwrap().unwrap();
+        assert!(loaded_thread.updated_at > original_updated);
+    }
+}