gitdataai/libs/api/chat/handlers/message.rs

use crate::ApiResponse;
use crate::error::ApiError;
use actix_web::{HttpResponse, Result, web};
use models::ai::AiMessage;
use sea_orm::EntityTrait;
use service::error::AppError;
use session::Session;
use uuid::Uuid;

use super::types::{CreateMessageParams, EditMessageParams, MessageListQuery, MessageResponse};

fn get_user_id(session: &Session) -> Result<Uuid, ApiError> {
    session
        .user()
        .ok_or_else(|| ApiError::from(AppError::Unauthorized))
}

#[utoipa::path(
    get,
    path = "/api/ai/conversations/{conversation_id}/messages",
    operation_id = "ai_message_list",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("limit" = Option<i64>, Query, description = "Max messages"),
    ),
    responses(
        (status = 200, description = "List messages", body = ApiResponse<Vec<MessageResponse>>),
        (status = 404, description = "Not found"),
    ),
    tag = "AI Chat"
)]
pub async fn message_list(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<Uuid>,
    query: web::Query<MessageListQuery>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let conversation_id = path.into_inner();

    let limit = query.limit.unwrap_or(50) as u64;
    let msgs = service
        .list_messages(conversation_id, user_id, limit)
        .await?;

    let resp: Vec<MessageResponse> = msgs.into_iter().map(MessageResponse::from).collect();

    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    post,
    path = "/api/ai/conversations/{conversation_id}/messages",
    operation_id = "ai_message_create",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
    ),
    request_body = CreateMessageParams,
    responses(
        (status = 200, description = "Message created", body = ApiResponse<MessageResponse>),
        (status = 404, description = "Not found"),
    ),
    tag = "AI Chat"
)]
pub async fn message_create(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<Uuid>,
    params: web::Json<CreateMessageParams>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let conversation_id = path.into_inner();

    let msg = service
        .create_message(
            conversation_id,
            user_id,
            params.parent_message_id,
            "user".to_string(),
            params.content.content.clone(),
            params.model.clone(),
            params.is_fork_origin.unwrap_or(false),
            params.metadata.clone(),
            params.room_id,
        )
        .await?;

    let resp = MessageResponse::from(msg);
    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    get,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}",
    operation_id = "ai_message_get",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID"),
    ),
    responses(
        (status = 200, description = "Get message", body = ApiResponse<MessageResponse>),
        (status = 404, description = "Not found"),
    ),
    tag = "AI Chat"
)]
pub async fn message_get(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    let msg = service
        .get_message(conversation_id, user_id, message_id)
        .await?;

    let resp = MessageResponse::from(msg);
    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    post,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/stop",
    operation_id = "ai_message_stop",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID"),
    ),
    responses(
        (status = 200, description = "Message stopped"),
    ),
    tag = "AI Chat"
)]
pub async fn message_stop(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    service
        .stop_message(conversation_id, user_id, message_id)
        .await?;

    Ok(crate::api_success())
}

#[utoipa::path(
    post,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/resend",
    operation_id = "ai_message_resend",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID"),
    ),
    responses(
        (status = 200, description = "Resend message", body = ApiResponse<MessageResponse>),
    ),
    tag = "AI Chat"
)]
pub async fn message_resend(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    let new_msg = service
        .resend_message(conversation_id, user_id, message_id)
        .await?;

    let resp = MessageResponse::from(new_msg);
    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    get,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/children",
    operation_id = "ai_message_children",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Parent message ID"),
    ),
    responses(
        (status = 200, description = "List child messages", body = ApiResponse<Vec<MessageResponse>>),
    ),
    tag = "AI Chat"
)]
pub async fn message_children(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, parent_message_id) = path.into_inner();

    let msgs = service
        .list_child_messages(conversation_id, user_id, parent_message_id)
        .await?;

    let resp: Vec<MessageResponse> = msgs.into_iter().map(MessageResponse::from).collect();

    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    get,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/stream",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID"),
    ),
    responses(
        (status = 200, description = "SSE stream"),
    ),
    tag = "AI Chat"
)]
pub async fn message_stream(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    // Streaming triggers AI execution and billing, so view-only access is not enough.
    let conv = service
        .find_conversation_full_access(conversation_id, user_id)
        .await?;

    let model = conv.model;

    let msg = AiMessage::find_by_id(message_id)
        .one(service.db.reader())
        .await
        .map_err(AppError::from)?
        .ok_or_else(|| ApiError::from(AppError::NotFound("message".into())))?;
    if msg.conversation_id != conversation_id || msg.role != "user" || !msg.is_latest {
        return Err(ApiError::from(AppError::NotFound("message".into())));
    }

    let response = actix_web::HttpResponse::Ok()
        .content_type("text/event-stream")
        .insert_header(("Cache-Control", "no-cache"))
        .insert_header(("X-Accel-Buffering", "no"))
        .streaming(super::super::stream::create_chat_sse_stream(
            service.get_ref().clone(),
            conversation_id,
            message_id,
            model,
            user_id,
        ));

    Ok(response.into())
}

#[utoipa::path(
    post,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/edit",
    operation_id = "ai_message_edit",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID to edit"),
    ),
    request_body = EditMessageParams,
    responses(
        (status = 200, description = "Message edited, new version created", body = ApiResponse<MessageResponse>),
    ),
    tag = "AI Chat"
)]
pub async fn message_edit(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
    params: web::Json<EditMessageParams>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    let new_msg = service
        .edit_message(conversation_id, user_id, message_id, params.content.clone())
        .await?;

    let resp = MessageResponse::from(new_msg);
    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    get,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/versions",
    operation_id = "ai_message_versions",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID"),
    ),
    responses(
        (status = 200, description = "List message versions", body = ApiResponse<Vec<MessageResponse>>),
    ),
    tag = "AI Chat"
)]
pub async fn message_versions(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    let versions = service
        .list_message_versions(conversation_id, user_id, message_id)
        .await?;

    let resp: Vec<MessageResponse> = versions.into_iter().map(MessageResponse::from).collect();

    Ok(ApiResponse::ok(resp).to_response())
}

#[utoipa::path(
    post,
    path = "/api/ai/conversations/{conversation_id}/messages/{message_id}/switch-version",
    operation_id = "ai_message_switch_version",
    params(
        ("conversation_id" = Uuid, Path, description = "Conversation ID"),
        ("message_id" = Uuid, Path, description = "Message ID"),
    ),
    request_body = super::types::SwitchVersionParams,
    responses(
        (status = 200, description = "Version switched", body = ApiResponse<MessageResponse>),
    ),
    tag = "AI Chat"
)]
pub async fn message_switch_version(
    service: web::Data<service::AppService>,
    session: Session,
    path: web::Path<(Uuid, Uuid)>,
    params: web::Json<super::types::SwitchVersionParams>,
) -> Result<HttpResponse, ApiError> {
    let user_id = get_user_id(&session)?;
    let (conversation_id, message_id) = path.into_inner();

    let msg = service
        .switch_message_version(conversation_id, user_id, message_id, params.version_number)
        .await?;

    let resp = MessageResponse::from(msg);
    Ok(ApiResponse::ok(resp).to_response())
}