gitdataai/libs/service/chat/fork.rs

use crate::error::AppError;
use models::ai::{AiMessage, ai_conversation, ai_message, ai_message_fork};
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, QueryFilter, QueryOrder, Set};
use uuid::Uuid;

use crate::AppService;

impl AppService {
    /// Fork a conversation from a specific message: creates a new conversation
    /// with all messages up to and including the source message, then copies
    /// the source message content as a new user message in the forked conversation.
    ///
    /// Returns the new conversation.
    pub async fn fork_conversation_from_message(
        &self,
        user_id: Uuid,
        conversation_id: Uuid,
        source_message_id: Uuid,
    ) -> Result<ai_conversation::Model, AppError> {
        let c = self
            .find_conversation_owned(conversation_id, user_id)
            .await?;

        // Verify source message exists in this conversation
        let source_msg = AiMessage::find_by_id(source_message_id)
            .one(self.db.reader())
            .await?
            .ok_or_else(|| AppError::NotFound("message".into()))?;

        if source_msg.conversation_id != conversation_id {
            return Err(AppError::NotFound("message not in conversation".into()));
        }

        // Get all messages in the conversation up to the source message
        let all_messages = AiMessage::find()
            .filter(ai_message::Column::ConversationId.eq(conversation_id))
            .filter(ai_message::Column::IsLatest.eq(true))
            .order_by_asc(ai_message::Column::CreatedAt)
            .all(self.db.reader())
            .await?;

        // Find the index of the source message in the ordered list
        let source_idx = all_messages
            .iter()
            .position(|m| m.id == source_message_id)
            .ok_or_else(|| AppError::NotFound("source message not found in conversation".into()))?;

        // Messages to copy: up to and including the source message
        let messages_to_copy = &all_messages[..=source_idx];

        // Create new conversation
        let new_conv_id = Uuid::new_v4();
        let now = chrono::Utc::now();
        let new_conv = ai_conversation::ActiveModel {
            id: Set(new_conv_id),
            user_id: Set(user_id),
            project_id: Set(c.project_id),
            scope: Set(c.scope.clone()),
            title: Set(c.title.clone().map(|t| format!("Fork: {}", t))),
            model: Set(c.model.clone()),
            model_config: Set(c.model_config.clone()),
            status: Set("active".to_string()),
            root_message_id: Set(None),
            fork_count: Set(0),
            is_shared: Set(false),
            message_count: Set(0),
            token_usage_total: Set(None),
            access_visibility: Set(c.access_visibility.clone()),
            can_ask: Set(c.can_ask.clone()),
            project_uid: Set(c.project_uid),
            model_uid: Set(c.model_uid),
            model_name: Set(c.model_name.clone()),
            created_at: Set(now),
            updated_at: Set(now),
        }
        .insert(self.db.writer())
        .await?;

        // Copy messages to the new conversation
        let mut prev_msg_id: Option<Uuid> = None;
        let mut msg_count = 0;

        for msg in messages_to_copy {
            let new_msg_id = Uuid::now_v7();
            ai_message::ActiveModel {
                id: Set(new_msg_id),
                conversation_id: Set(new_conv_id),
                parent_message_id: Set(prev_msg_id),
                role: Set(msg.role.clone()),
                content: Set(msg.content.clone()),
                model: Set(msg.model.clone()),
                is_fork_origin: Set(false),
                stop_reason: Set(None),
                input_tokens: Set(None),
                output_tokens: Set(None),
                latency_ms: Set(None),
                metadata: Set(None),
                room_id: Set(None),
                version_group_id: Set(Some(new_msg_id)),
                version_number: Set(1),
                is_latest: Set(true),
                created_at: Set(chrono::Utc::now()),
            }
            .insert(self.db.writer())
            .await?;

            prev_msg_id = Some(new_msg_id);
            msg_count += 1;
        }

        // Update conversation message count
        let mut updated: ai_conversation::ActiveModel = new_conv.clone().into();
        updated.message_count = Set(msg_count);
        updated.updated_at = Set(chrono::Utc::now());
        updated.update(self.db.writer()).await?;

        // Mark source message as fork origin
        let mut source: ai_message::ActiveModel = source_msg.into();
        source.is_fork_origin = Set(true);
        source.update(self.db.writer()).await?;

        // Create fork record
        ai_message_fork::ActiveModel {
            id: Set(Uuid::new_v4()),
            conversation_id: Set(Some(new_conv_id)),
            source_message_id: Set(source_message_id),
            fork_message_id: Set(prev_msg_id.unwrap_or(source_message_id)),
            created_at: Set(chrono::Utc::now()),
        }
        .insert(self.db.writer())
        .await?;

        // Update original conversation fork_count
        let original_fork_count = c.fork_count;
        let mut orig_updated: ai_conversation::ActiveModel = c.into();
        orig_updated.fork_count = Set(original_fork_count + 1);
        orig_updated.updated_at = Set(chrono::Utc::now());
        orig_updated.update(self.db.writer()).await?;

        // Return the new conversation
        let new_conv = ai_conversation::Entity::find_by_id(new_conv_id)
            .one(self.db.reader())
            .await?
            .ok_or_else(|| AppError::NotFound("conversation".into()))?;

        Ok(new_conv)
    }

    /// List all fork records for a message within a conversation.
    pub async fn list_forks(
        &self,
        conversation_id: Uuid,
        user_id: Uuid,
        source_message_id: Uuid,
    ) -> Result<Vec<ai_message_fork::Model>, AppError> {
        self.find_conversation_owned(conversation_id, user_id)
            .await?;

        let forks = ai_message_fork::Entity::find()
            .filter(ai_message_fork::Column::SourceMessageId.eq(source_message_id))
            .all(self.db.reader())
            .await?;

        Ok(forks)
    }
}