gitdataai/libs/room/src/message.rs

use crate::error::RoomError;
use crate::service::RoomService;
use crate::ws_context::WsUserContext;
use chrono::Utc;
use models::rooms::{room, room_message, room_thread};
use models::users::user as user_model;
use queue::RoomMessageEnvelope;
use sea_orm::*;
use serde_json;
use uuid::Uuid;

impl RoomService {
    pub async fn room_message_list(
        &self,
        room_id: Uuid,
        before_seq: Option<i64>,
        after_seq: Option<i64>,
        limit: Option<u64>,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageListResponse, RoomError> {
        let user_id = ctx.user_id;
        self.require_room_member(room_id, user_id).await?;

        let mut query = room_message::Entity::find().filter(room_message::Column::Room.eq(room_id));
        if let Some(before_seq) = before_seq {
            query = query.filter(room_message::Column::Seq.lt(before_seq));
        }
        if let Some(after_seq) = after_seq {
            query = query.filter(room_message::Column::Seq.gt(after_seq));
        }

        let total = query.clone().count(&self.db).await? as i64;
        let models = query
            .order_by_desc(room_message::Column::Seq)
            .limit(limit.unwrap_or(50))
            .all(&self.db)
            .await?;

        let user_ids: Vec<Uuid> = models
            .iter()
            .filter(|m| m.sender_type.to_string() == "member")
            .filter_map(|m| m.sender_id)
            .collect();
        let ai_model_ids: Vec<Uuid> = models
            .iter()
            .filter(|m| m.sender_type.to_string() == "ai")
            .filter_map(|m| m.sender_id)
            .collect();

        let users: std::collections::HashMap<Uuid, String> = if !user_ids.is_empty() {
            use sea_orm::ColumnTrait;
            user_model::Entity::find()
                .filter(user_model::Column::Uid.is_in(user_ids))
                .all(&self.db)
                .await?
                .into_iter()
                .map(|u| (u.uid, u.display_name.unwrap_or(u.username)))
                .collect()
        } else {
            std::collections::HashMap::new()
        };

        let ai_names: std::collections::HashMap<Uuid, String> = if !ai_model_ids.is_empty() {
            use sea_orm::ColumnTrait;
            models::agents::model::Entity::find()
                .filter(models::agents::model::Column::Id.is_in(ai_model_ids))
                .all(&self.db)
                .await?
                .into_iter()
                .map(|m| (m.id, m.name))
                .collect()
        } else {
            std::collections::HashMap::new()
        };

        let mut messages: Vec<super::RoomMessageResponse> = models
            .into_iter()
            .map(|msg| {
                let sender_type = msg.sender_type.to_string();
                let display_name = match sender_type.as_str() {
                    "ai" => msg.sender_id.and_then(|id| ai_names.get(&id).cloned()),
                    _ => msg.sender_id.and_then(|id| users.get(&id).cloned()),
                };
                super::RoomMessageResponse {
                    id: msg.id,
                    seq: msg.seq,
                    room: msg.room,
                    sender_type,
                    sender_id: msg.sender_id,
                    display_name,
                    thread: msg.thread,
                    in_reply_to: msg.in_reply_to,
                    content: msg.content,
                    content_type: msg.content_type.to_string(),
                    edited_at: msg.edited_at,
                    send_at: msg.send_at,
                    revoked: msg.revoked,
                    revoked_by: msg.revoked_by,
                }
            })
            .collect();
        messages.reverse();

        Ok(super::RoomMessageListResponse { messages, total })
    }

    pub async fn room_message_create(
        &self,
        room_id: Uuid,
        request: super::RoomMessageCreateRequest,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageResponse, RoomError> {
        let user_id = ctx.user_id;
        let room_model = self.find_room_or_404(room_id).await?;
        self.require_room_member(room_id, user_id).await?;

        let content_type_str = request
            .content_type
            .clone()
            .unwrap_or_else(|| "text".to_string());
        Self::parse_message_content_type(Some(content_type_str.clone()))?;
        Self::validate_content(&request.content, super::MAX_MESSAGE_CONTENT_LEN)?;
        let content = Self::sanitize_content(&request.content);

        let thread_id = request.thread;

        if let Some(tid) = thread_id {
            let thread = room_thread::Entity::find_by_id(tid)
                .one(&self.db)
                .await?
                .ok_or_else(|| RoomError::NotFound("Thread not found".to_string()))?;
            if thread.room != room_id {
                return Err(RoomError::BadRequest("thread not in room".to_string()));
            }
        }

        let seq = self.next_room_message_seq(room_id, &self.db).await?;
        let now = Utc::now();
        let id = Uuid::now_v7();
        let project_id = room_model.project;

        let in_reply_to = request.in_reply_to;
        let envelope = RoomMessageEnvelope {
            id,
            dedup_key: Some(format!("{}:{}", room_id, id)),
            room_id,
            sender_type: "member".to_string(),
            sender_id: Some(user_id),
            thread_id,
            in_reply_to,
            content: content.clone(),
            content_type: content_type_str.clone(),
            send_at: now,
            seq,
        };

        let db = &self.db;
        let txn = db.begin().await?;

        self.queue.publish(room_id, envelope).await?;
        self.room_manager.metrics.messages_sent.increment(1);

        let mut room_active: room::ActiveModel = room_model.clone().into();
        room_active.last_msg_at = Set(now);
        room_active.update(&txn).await?;

        if let Some(tid) = thread_id {
            let thread = room_thread::Entity::find_by_id(tid)
                .one(&txn)
                .await?
                .ok_or_else(|| RoomError::NotFound("Thread not found".to_string()))?;

            let participants: Vec<Uuid> =
                serde_json::from_value(thread.participants.clone()).unwrap_or_default();
            let participants: Vec<Uuid> = if !participants.contains(&user_id) {
                let mut p = participants;
                p.push(user_id);
                p
            } else {
                participants
            };

            let preview = if content.len() > 50 {
                format!("{}...", &content[..50])
            } else {
                content.clone()
            };

            let mut active: room_thread::ActiveModel = thread.into();
            active.last_message_at = Set(now);
            active.last_message_preview = Set(Some(preview));
            active.participants = Set(serde_json::to_value(participants).unwrap_or_default());
            active.updated_at = Set(now);
            active.update(&txn).await?;
        }

        txn.commit().await?;

        self.publish_room_event(
            project_id,
            super::RoomEventType::NewMessage,
            Some(room_id),
            None,
            Some(id),
            Some(seq),
        )
        .await;

        let mentioned_users = self.resolve_mentions(&request.content).await;
        for mentioned_user_id in mentioned_users {
            if mentioned_user_id == user_id {
                continue;
            }
            let _ = self
                .notification_create(super::NotificationCreateRequest {
                    notification_type: super::NotificationType::Mention,
                    user_id: mentioned_user_id,
                    title: format!("{} 在 {} 中提到了你", user_id, room_model.room_name),
                    content: Some(content.clone()),
                    room_id: Some(room_id),
                    project_id,
                    related_message_id: Some(id),
                    related_user_id: Some(user_id),
                    related_room_id: Some(room_id),
                    metadata: None,
                    expires_at: None,
                })
                .await;
        }

        let should_respond = self.should_ai_respond(room_id).await.unwrap_or(false);
        let is_text_message = request
            .content_type
            .as_ref()
            .map(|ct| ct == "text")
            .unwrap_or(true);
        if should_respond && is_text_message {
            if let Err(e) = self
                .process_message_ai(room_id, id, user_id, content.clone())
                .await
            {
                slog::warn!(self.log, "Failed to process AI message: {}", e);
            }
        }

        let display_name = {
            let user = user_model::Entity::find()
                .filter(user_model::Column::Uid.eq(user_id))
                .one(&self.db)
                .await
                .ok()
                .flatten();
            user.map(|u| u.display_name.unwrap_or_else(|| u.username))
        };

        Ok(super::RoomMessageResponse {
            id,
            seq,
            room: room_id,
            sender_type: "member".to_string(),
            sender_id: Some(user_id),
            display_name,
            thread: thread_id,
            in_reply_to,
            content: request.content,
            content_type: content_type_str,
            edited_at: None,
            send_at: now,
            revoked: None,
            revoked_by: None,
        })
    }

    pub async fn room_message_update(
        &self,
        message_id: Uuid,
        request: super::RoomMessageUpdateRequest,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageResponse, RoomError> {
        let user_id = ctx.user_id;
        Self::validate_content(&request.content, super::MAX_MESSAGE_CONTENT_LEN)?;
        let model = room_message::Entity::find_by_id(message_id)
            .one(&self.db)
            .await?
            .ok_or_else(|| RoomError::NotFound("Message not found".to_string()))?;

        self.require_room_member(model.room, user_id).await?;
        if model.sender_id != Some(user_id) {
            return Err(RoomError::NoPower);
        }

        let elapsed = Utc::now().signed_duration_since(model.send_at);
        if elapsed.num_minutes() > 120 {
            return Err(RoomError::BadRequest(
                "消息只能在发送后 2 小时内编辑".into(),
            ));
        }

        let old_content = model.content.clone();
        let new_content = request.content.clone();
        self.save_message_edit_history(message_id, user_id, old_content, new_content)
            .await?;

        let mut active: room_message::ActiveModel = model.into();
        active.content = Set(request.content);
        active.edited_at = Set(Some(Utc::now()));
        let updated = active.update(&self.db).await?;
        let updated_room = updated.room;

        let room = self.find_room_or_404(updated_room).await?;
        self.publish_room_event(
            room.project,
            super::RoomEventType::MessageEdited,
            Some(updated_room),
            None,
            Some(updated.id),
            None,
        )
        .await;

        Ok(self.resolve_display_name(updated, updated_room).await)
    }

    pub async fn room_message_revoke(
        &self,
        message_id: Uuid,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageResponse, RoomError> {
        let user_id = ctx.user_id;
        let model = room_message::Entity::find_by_id(message_id)
            .one(&self.db)
            .await?
            .ok_or_else(|| RoomError::NotFound("Message not found".to_string()))?;

        let member = self.require_room_member_model(model.room, user_id).await?;
        let can_admin = Self::is_room_admin(&member.role);
        let can_author = model.sender_id == Some(user_id);
        if !can_admin && !can_author {
            return Err(RoomError::NoPower);
        }

        let mut active: room_message::ActiveModel = model.into();
        active.revoked = Set(Some(Utc::now()));
        active.revoked_by = Set(Some(user_id));
        let updated = active.update(&self.db).await?;
        let updated_room = updated.room;

        let room = self.find_room_or_404(updated_room).await?;
        self.publish_room_event(
            room.project,
            super::RoomEventType::MessageRevoked,
            Some(updated_room),
            None,
            Some(updated.id),
            None,
        )
        .await;

        Ok(self.resolve_display_name(updated, updated_room).await)
    }

    pub async fn room_message_get(
        &self,
        message_id: Uuid,
        ctx: &WsUserContext,
    ) -> Result<super::RoomMessageResponse, RoomError> {
        let user_id = ctx.user_id;
        let model = room_message::Entity::find_by_id(message_id)
            .one(&self.db)
            .await?
            .ok_or_else(|| RoomError::NotFound("Message not found".to_string()))?;
        let room_id = model.room;
        self.require_room_member(room_id, user_id).await?;
        Ok(self.resolve_display_name(model, room_id).await)
    }
}