enact-core 0.0.2

//! 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);
    }
}