feat(room): store ordered streaming chunks + billing integration
- Save thinking_content as {"__chunks__": [{type, content}]} for replay
- Tool call sanitization — don't expose raw results to frontend
- Billing record_ai_usage integration
- Room service module refactoring into service/ directory
This commit is contained in:
parent
b4b5538447
commit
f5e3da35b0
@ -188,6 +188,7 @@ impl MessageProducer {
|
||||
in_reply_to: None,
|
||||
content: String::new(),
|
||||
content_type: String::new(),
|
||||
thinking_content: None,
|
||||
send_at: chrono::Utc::now(),
|
||||
seq: 0,
|
||||
display_name: None,
|
||||
|
||||
@ -17,6 +17,9 @@ pub struct RoomMessageEnvelope {
|
||||
pub in_reply_to: Option<Uuid>,
|
||||
pub content: String,
|
||||
pub content_type: String,
|
||||
/// Accumulated AI reasoning/thinking text.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking_content: Option<String>,
|
||||
pub send_at: DateTime<Utc>,
|
||||
pub seq: i64,
|
||||
/// Pre-resolved display name for the sender (e.g. AI model name).
|
||||
@ -34,6 +37,9 @@ pub struct RoomMessageEvent {
|
||||
pub in_reply_to: Option<Uuid>,
|
||||
pub content: String,
|
||||
pub content_type: String,
|
||||
/// Accumulated AI reasoning/thinking text.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking_content: Option<String>,
|
||||
pub send_at: DateTime<Utc>,
|
||||
pub seq: i64,
|
||||
pub display_name: Option<String>,
|
||||
@ -79,6 +85,7 @@ impl From<RoomMessageEnvelope> for RoomMessageEvent {
|
||||
in_reply_to: e.in_reply_to,
|
||||
content: e.content,
|
||||
content_type: e.content_type,
|
||||
thinking_content: e.thinking_content,
|
||||
send_at: e.send_at,
|
||||
seq: e.seq,
|
||||
display_name: e.display_name,
|
||||
|
||||
@ -826,6 +826,7 @@ pub fn make_persist_fn(
|
||||
thread: Set(env.thread_id),
|
||||
content: Set(env.content.clone()),
|
||||
content_type: Set(content_type),
|
||||
thinking_content: Set(env.thinking_content.clone()),
|
||||
edited_at: Set(None),
|
||||
send_at: Set(env.send_at.clone()),
|
||||
revoked: Set(None),
|
||||
|
||||
@ -32,3 +32,9 @@ impl From<anyhow::Error> for RoomError {
|
||||
RoomError::Internal(e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<agent::error::AgentError> for RoomError {
|
||||
fn from(e: agent::error::AgentError) -> Self {
|
||||
RoomError::Internal(e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
@ -68,6 +68,7 @@ impl From<room_message::Model> for super::RoomMessageResponse {
|
||||
thread: value.thread,
|
||||
content: value.content,
|
||||
content_type: value.content_type.to_string(),
|
||||
thinking_content: value.thinking_content,
|
||||
edited_at: value.edited_at,
|
||||
send_at: value.send_at,
|
||||
revoked: value.revoked,
|
||||
@ -427,6 +428,7 @@ impl RoomService {
|
||||
thread: msg.thread,
|
||||
content: msg.content,
|
||||
content_type: msg.content_type.to_string(),
|
||||
thinking_content: msg.thinking_content,
|
||||
edited_at: msg.edited_at,
|
||||
send_at: msg.send_at,
|
||||
revoked: msg.revoked,
|
||||
|
||||
@ -92,6 +92,7 @@ impl RoomService {
|
||||
in_reply_to: msg.in_reply_to,
|
||||
content: msg.content,
|
||||
content_type: msg.content_type.to_string(),
|
||||
thinking_content: msg.thinking_content,
|
||||
edited_at: msg.edited_at,
|
||||
send_at: msg.send_at,
|
||||
revoked: msg.revoked,
|
||||
@ -158,7 +159,7 @@ impl RoomService {
|
||||
}
|
||||
}
|
||||
|
||||
let seq = Self::next_room_message_seq_internal(room_id, &self.db, &self.cache).await?;
|
||||
let seq = crate::service::next_room_message_seq_internal(room_id, &self.db, &self.cache).await?;
|
||||
let now = Utc::now();
|
||||
let id = Uuid::now_v7();
|
||||
let project_id = room_model.project;
|
||||
@ -175,6 +176,7 @@ impl RoomService {
|
||||
in_reply_to,
|
||||
content: content.clone(),
|
||||
content_type: content_type_str.clone(),
|
||||
thinking_content: None,
|
||||
send_at: now,
|
||||
seq,
|
||||
display_name: None,
|
||||
@ -349,6 +351,7 @@ impl RoomService {
|
||||
in_reply_to,
|
||||
content: request.content,
|
||||
content_type: content_type_str,
|
||||
thinking_content: None,
|
||||
edited_at: None,
|
||||
send_at: now,
|
||||
revoked: None,
|
||||
|
||||
@ -321,6 +321,7 @@ impl RoomService {
|
||||
in_reply_to: msg.in_reply_to,
|
||||
content: msg.content,
|
||||
content_type: msg.content_type.to_string(),
|
||||
thinking_content: msg.thinking_content,
|
||||
edited_at: msg.edited_at,
|
||||
send_at: msg.send_at,
|
||||
revoked: msg.revoked,
|
||||
|
||||
@ -124,6 +124,7 @@ impl RoomService {
|
||||
in_reply_to: row.try_get::<Option<MessageId>>("", "in_reply_to").ok().flatten(),
|
||||
content: row.try_get::<String>("", "content").unwrap_or_default(),
|
||||
content_type,
|
||||
thinking_content: None,
|
||||
edited_at: row.try_get::<Option<DateTimeUtc>>("", "edited_at").ok().flatten(),
|
||||
send_at: row.try_get::<DateTimeUtc>("", "send_at").unwrap_or_default(),
|
||||
revoked: row.try_get::<Option<DateTimeUtc>>("", "revoked").ok().flatten(),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
73
libs/room/src/service/access.rs
Normal file
73
libs/room/src/service/access.rs
Normal file
@ -0,0 +1,73 @@
|
||||
use db::database::AppDatabase;
|
||||
use models::projects::project_members;
|
||||
use models::rooms::room;
|
||||
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::RoomError;
|
||||
|
||||
pub async fn check_room_access(
|
||||
db: &AppDatabase,
|
||||
room_id: Uuid,
|
||||
user_id: Uuid,
|
||||
) -> Result<(), RoomError> {
|
||||
let room = room::Entity::find_by_id(room_id)
|
||||
.one(db)
|
||||
.await?
|
||||
.ok_or_else(|| RoomError::NotFound("Room not found".to_string()))?;
|
||||
|
||||
if room.public {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if require_room_member(db, room_id, user_id).await.is_ok() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
check_project_member(db, room.project, user_id).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn check_project_member(
|
||||
db: &AppDatabase,
|
||||
project_id: Uuid,
|
||||
user_id: Uuid,
|
||||
) -> Result<(), RoomError> {
|
||||
let member = project_members::Entity::find()
|
||||
.filter(project_members::Column::Project.eq(project_id))
|
||||
.filter(project_members::Column::User.eq(user_id))
|
||||
.one(db)
|
||||
.await?;
|
||||
|
||||
if member.is_some() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(RoomError::NoPower)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn require_room_member(
|
||||
db: &AppDatabase,
|
||||
room_id: Uuid,
|
||||
user_id: Uuid,
|
||||
) -> Result<(), RoomError> {
|
||||
use models::rooms::room_member::{Column as RmCol, Entity as RoomMember};
|
||||
|
||||
let member = RoomMember::find()
|
||||
.filter(RmCol::Room.eq(room_id))
|
||||
.filter(RmCol::User.eq(user_id))
|
||||
.one(db)
|
||||
.await?;
|
||||
|
||||
member
|
||||
.ok_or_else(|| RoomError::NotFound("Room member not found".to_string()))
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub async fn find_room_or_404(db: &AppDatabase, room_id: Uuid) -> Result<room::Model, RoomError> {
|
||||
room::Entity::find_by_id(room_id)
|
||||
.one(db)
|
||||
.await?
|
||||
.ok_or_else(|| RoomError::NotFound("Room not found".to_string()))
|
||||
}
|
||||
77
libs/room/src/service/ai_common.rs
Normal file
77
libs/room/src/service/ai_common.rs
Normal file
@ -0,0 +1,77 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use queue::MessageProducer;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::sequence::next_room_message_seq_internal;
|
||||
use crate::connection::RoomConnectionManager;
|
||||
use crate::error::RoomError;
|
||||
|
||||
pub async fn create_and_publish_ai_message(
|
||||
db: &AppDatabase,
|
||||
cache: &AppCache,
|
||||
queue: &MessageProducer,
|
||||
room_manager: &Arc<RoomConnectionManager>,
|
||||
room_id: Uuid,
|
||||
project_id: Uuid,
|
||||
_reply_to: Uuid,
|
||||
content: String,
|
||||
model_id: Uuid,
|
||||
model_display_name: Option<String>,
|
||||
) -> Result<Uuid, RoomError> {
|
||||
let now = Utc::now();
|
||||
let seq = next_room_message_seq_internal(room_id, db, cache).await?;
|
||||
let id = Uuid::now_v7();
|
||||
|
||||
let envelope = queue::RoomMessageEnvelope {
|
||||
id,
|
||||
dedup_key: Some(format!("{}:{}", room_id, id)),
|
||||
room_id,
|
||||
sender_type: "ai".to_string(),
|
||||
sender_id: None,
|
||||
model_id: Some(model_id),
|
||||
thread_id: None,
|
||||
content: content.clone(),
|
||||
content_type: "text".to_string(),
|
||||
thinking_content: None,
|
||||
send_at: now,
|
||||
seq,
|
||||
in_reply_to: None,
|
||||
display_name: model_display_name.clone(),
|
||||
};
|
||||
|
||||
queue.publish(room_id, envelope).await?;
|
||||
room_manager.metrics.messages_sent.increment(1);
|
||||
|
||||
let event = queue::RoomMessageEvent {
|
||||
id,
|
||||
room_id,
|
||||
sender_type: "ai".to_string(),
|
||||
sender_id: None,
|
||||
thread_id: None,
|
||||
content: content.clone(),
|
||||
content_type: "text".to_string(),
|
||||
thinking_content: None,
|
||||
send_at: now,
|
||||
seq,
|
||||
display_name: model_display_name,
|
||||
in_reply_to: None,
|
||||
reactions: None,
|
||||
message_id: None,
|
||||
};
|
||||
room_manager.broadcast(room_id, event).await;
|
||||
|
||||
super::notifications::publish_room_event(
|
||||
queue,
|
||||
project_id,
|
||||
crate::RoomEventType::NewMessage,
|
||||
Some(room_id),
|
||||
Some(id),
|
||||
Some(seq),
|
||||
);
|
||||
|
||||
Ok(id)
|
||||
}
|
||||
94
libs/room/src/service/ai_nonstreaming.rs
Normal file
94
libs/room/src/service/ai_nonstreaming.rs
Normal file
@ -0,0 +1,94 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room_ai;
|
||||
use queue::MessageProducer;
|
||||
use sea_orm::{sea_query::Expr, ColumnTrait, EntityTrait, ExprTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::ai_common::create_and_publish_ai_message;
|
||||
use crate::connection::RoomConnectionManager;
|
||||
use agent::chat::{AiRequest, ChatService};
|
||||
|
||||
pub async fn process_message_ai_nonstreaming(
|
||||
chat_service: Arc<ChatService>,
|
||||
request: AiRequest,
|
||||
room_id: Uuid,
|
||||
project_id: Uuid,
|
||||
model_id: Uuid,
|
||||
lock_guard: crate::room_ai_queue::RoomAiLockGuard,
|
||||
db: AppDatabase,
|
||||
cache: AppCache,
|
||||
queue: MessageProducer,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
) {
|
||||
let chat_service = chat_service.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _lock_guard = lock_guard;
|
||||
let model_display_name = request.model.name.clone();
|
||||
match chat_service.process(request).await {
|
||||
Ok(result) => {
|
||||
if let Err(e) = create_and_publish_ai_message(
|
||||
&db,
|
||||
&cache,
|
||||
&queue,
|
||||
&room_manager,
|
||||
room_id,
|
||||
project_id,
|
||||
Uuid::now_v7(),
|
||||
result.content,
|
||||
model_id,
|
||||
Some(model_display_name),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(error = %e, "Failed to create AI message");
|
||||
} else {
|
||||
let now = Utc::now();
|
||||
if let Err(e) = room_ai::Entity::update_many()
|
||||
.col_expr(
|
||||
room_ai::Column::CallCount,
|
||||
Expr::col(room_ai::Column::CallCount).add(1),
|
||||
)
|
||||
.col_expr(room_ai::Column::LastCallAt, Expr::value(Some(now)))
|
||||
.filter(room_ai::Column::Room.eq(room_id))
|
||||
.filter(room_ai::Column::Model.eq(model_id))
|
||||
.exec(&db)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "Failed to update room_ai call stats");
|
||||
}
|
||||
|
||||
// Record billing (non-fatal)
|
||||
let _ = super::billing::record_ai_usage(
|
||||
&db,
|
||||
project_id,
|
||||
model_id,
|
||||
result.input_tokens,
|
||||
result.output_tokens,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "AI processing failed");
|
||||
let _ = create_and_publish_ai_message(
|
||||
&db,
|
||||
&cache,
|
||||
&queue,
|
||||
&room_manager,
|
||||
room_id,
|
||||
project_id,
|
||||
Uuid::now_v7(),
|
||||
format!("[AI error: {}]", e),
|
||||
model_id,
|
||||
Some(model_display_name),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
98
libs/room/src/service/ai_react_nonstreaming.rs
Normal file
98
libs/room/src/service/ai_react_nonstreaming.rs
Normal file
@ -0,0 +1,98 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room_ai;
|
||||
use queue::MessageProducer;
|
||||
use sea_orm::{sea_query::Expr, ColumnTrait, EntityTrait, ExprTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::ai_common::create_and_publish_ai_message;
|
||||
use crate::connection::RoomConnectionManager;
|
||||
use agent::chat::{AiRequest, ChatService};
|
||||
|
||||
pub async fn process_message_ai_react_nonstreaming(
|
||||
chat_service: Arc<ChatService>,
|
||||
request: AiRequest,
|
||||
room_id: Uuid,
|
||||
project_id: Uuid,
|
||||
model_id: Uuid,
|
||||
lock_guard: crate::room_ai_queue::RoomAiLockGuard,
|
||||
db: AppDatabase,
|
||||
cache: AppCache,
|
||||
queue: MessageProducer,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
let _lock_guard = lock_guard;
|
||||
let model_display_name = request.model.name.clone();
|
||||
|
||||
let final_answer = chat_service
|
||||
.process_react(&request, |_step| {})
|
||||
.await;
|
||||
|
||||
match final_answer {
|
||||
Ok(response) => {
|
||||
if let Err(e) = create_and_publish_ai_message(
|
||||
&db,
|
||||
&cache,
|
||||
&queue,
|
||||
&room_manager,
|
||||
room_id,
|
||||
project_id,
|
||||
Uuid::now_v7(),
|
||||
response,
|
||||
model_id,
|
||||
Some(model_display_name),
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::error!(error = %e, "Failed to create ReAct AI message");
|
||||
} else {
|
||||
let now = Utc::now();
|
||||
if let Err(e) = room_ai::Entity::update_many()
|
||||
.col_expr(
|
||||
room_ai::Column::CallCount,
|
||||
Expr::col(room_ai::Column::CallCount).add(1),
|
||||
)
|
||||
.col_expr(room_ai::Column::LastCallAt, Expr::value(Some(now)))
|
||||
.filter(room_ai::Column::Room.eq(room_id))
|
||||
.filter(room_ai::Column::Model.eq(model_id))
|
||||
.exec(&db)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "Failed to update room_ai call stats");
|
||||
}
|
||||
|
||||
// Record billing (non-fatal)
|
||||
// TODO: ReAct agent does not track token counts yet; billing with 0/0
|
||||
let _ = super::billing::record_ai_usage(
|
||||
&db,
|
||||
project_id,
|
||||
model_id,
|
||||
0,
|
||||
0,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "ReAct agent failed");
|
||||
let _ = create_and_publish_ai_message(
|
||||
&db,
|
||||
&cache,
|
||||
&queue,
|
||||
&room_manager,
|
||||
room_id,
|
||||
project_id,
|
||||
Uuid::now_v7(),
|
||||
format!("[AI error: {}]", e),
|
||||
model_id,
|
||||
Some(model_display_name),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
266
libs/room/src/service/ai_react_streaming.rs
Normal file
266
libs/room/src/service/ai_react_streaming.rs
Normal file
@ -0,0 +1,266 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room_ai;
|
||||
use queue::{MessageProducer, ProjectRoomEvent, RoomMessageEnvelope};
|
||||
use sea_orm::{sea_query::Expr, ColumnTrait, EntityTrait, ExprTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::sequence::next_room_message_seq_internal;
|
||||
use crate::connection::RoomConnectionManager;
|
||||
use agent::chat::{AiRequest, ChatService};
|
||||
use agent::react::ReactStep;
|
||||
|
||||
pub async fn process_message_ai_react_streaming(
|
||||
chat_service: Arc<ChatService>,
|
||||
request: AiRequest,
|
||||
room_id: Uuid,
|
||||
project_id: Uuid,
|
||||
model_id: Uuid,
|
||||
lock_guard: crate::room_ai_queue::RoomAiLockGuard,
|
||||
db: AppDatabase,
|
||||
_cache: AppCache,
|
||||
queue: MessageProducer,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
) {
|
||||
use queue::RoomMessageStreamChunkEvent;
|
||||
|
||||
let streaming_msg_id = Uuid::now_v7();
|
||||
let seq = match next_room_message_seq_internal(room_id, &db, &_cache).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "Failed to get seq for ReAct streaming");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let room_id_inner = room_id;
|
||||
let project_id_inner = project_id;
|
||||
let now = Utc::now();
|
||||
let sender_type = "ai".to_string();
|
||||
let ai_display_name = request.model.name.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _lock_guard = lock_guard;
|
||||
|
||||
// Collect ordered steps for storage and streaming.
|
||||
let steps: std::sync::Arc<std::sync::Mutex<Vec<(String, String)>>> =
|
||||
std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
|
||||
let last_action_name: std::sync::Arc<std::sync::Mutex<String>> =
|
||||
std::sync::Arc::new(std::sync::Mutex::new(String::new()));
|
||||
let answer_buffer: std::sync::Arc<std::sync::Mutex<String>> =
|
||||
std::sync::Arc::new(std::sync::Mutex::new(String::new()));
|
||||
let step_count = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
|
||||
let on_step = {
|
||||
let room_manager = room_manager.clone();
|
||||
let streaming_msg_id = streaming_msg_id;
|
||||
let room_id = room_id_inner;
|
||||
let step_count = step_count.clone();
|
||||
let ai_display_name_for_step = std::sync::Arc::new(ai_display_name.clone());
|
||||
let steps = steps.clone();
|
||||
let answer_buffer = answer_buffer.clone();
|
||||
let last_action_name = last_action_name.clone();
|
||||
move |step: ReactStep| {
|
||||
let room_manager = room_manager.clone();
|
||||
let (chunk_type, content) = match &step {
|
||||
ReactStep::Thought { step: _, thought } => {
|
||||
("thinking".to_string(), format!("[Thinking] {}", thought))
|
||||
}
|
||||
ReactStep::Action { step: _, action } => {
|
||||
*last_action_name.lock().unwrap() = action.name.clone();
|
||||
("tool_call".to_string(), format!("[Action] Calling `{}` with {:?}", action.name, action.args))
|
||||
}
|
||||
ReactStep::Observation {
|
||||
step: _,
|
||||
observation: _,
|
||||
} => {
|
||||
// Sanitize observation — don't expose raw tool output to frontend
|
||||
let action_name = last_action_name.lock().unwrap().clone();
|
||||
("tool_call".to_string(), format!("[Observation] {} (completed)", action_name))
|
||||
}
|
||||
ReactStep::Answer { step: _, answer } => {
|
||||
("answer".to_string(), answer.clone())
|
||||
}
|
||||
};
|
||||
|
||||
let is_answer = matches!(&step, ReactStep::Answer { .. });
|
||||
if is_answer {
|
||||
step_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// Record ordered step for storage
|
||||
{
|
||||
let mut s = steps.lock().unwrap();
|
||||
s.push((chunk_type.clone(), content.clone()));
|
||||
}
|
||||
if is_answer {
|
||||
let mut ab = answer_buffer.lock().unwrap();
|
||||
ab.push_str(&content);
|
||||
}
|
||||
|
||||
let done = is_answer;
|
||||
let ai_name = ai_display_name_for_step.clone();
|
||||
tokio::spawn(async move {
|
||||
let event = RoomMessageStreamChunkEvent {
|
||||
message_id: streaming_msg_id,
|
||||
room_id,
|
||||
content: content.clone(),
|
||||
done,
|
||||
error: None,
|
||||
display_name: Some((*ai_name).clone()),
|
||||
chunk_type: Some(chunk_type),
|
||||
};
|
||||
room_manager.broadcast_stream_chunk(event).await;
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let result = chat_service.process_react(&request, on_step).await;
|
||||
|
||||
let final_content = answer_buffer.lock().unwrap().clone();
|
||||
let all_steps = steps.lock().unwrap().clone();
|
||||
let reasoning_chain: String = all_steps
|
||||
.iter()
|
||||
.filter(|(t, _)| t != "answer")
|
||||
.map(|(_, c)| c.clone())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
let content_to_persist = if !final_content.is_empty() {
|
||||
final_content
|
||||
} else if !reasoning_chain.trim().is_empty() {
|
||||
format!(
|
||||
"[Agent ran through {} reasoning steps but did not produce a final answer.]\n{}",
|
||||
step_count.load(std::sync::atomic::Ordering::Relaxed),
|
||||
reasoning_chain.trim_end()
|
||||
)
|
||||
} else {
|
||||
String::from("[No output from reasoning agent]")
|
||||
};
|
||||
|
||||
let (err_msg, should_log) = match &result {
|
||||
Err(e) => (Some(format!("[Agent error: {}]", e)), true),
|
||||
_ => (None, false),
|
||||
};
|
||||
|
||||
let content_to_persist = if let Some(msg) = &err_msg {
|
||||
format!(
|
||||
"{}\n[Error during reasoning: {}]",
|
||||
content_to_persist.trim_end(),
|
||||
msg.trim_start_matches("[Agent error: ")
|
||||
.trim_end_matches("]")
|
||||
)
|
||||
} else {
|
||||
content_to_persist
|
||||
};
|
||||
|
||||
if should_log {
|
||||
tracing::error!(error = %result.as_ref().unwrap_err(), "ReAct streaming failed");
|
||||
}
|
||||
|
||||
let persist_content = content_to_persist.trim().to_string();
|
||||
if persist_content.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Serialize ordered steps as JSON for ordered replay.
|
||||
let thinking_content = {
|
||||
let steps = steps.lock().unwrap();
|
||||
if steps.is_empty() {
|
||||
None
|
||||
} else {
|
||||
let chunks_json = serde_json::json!({
|
||||
"__chunks__": steps.iter().map(|(t, c)| serde_json::json!({
|
||||
"type": t,
|
||||
"content": c,
|
||||
})).collect::<Vec<_>>(),
|
||||
});
|
||||
Some(chunks_json.to_string())
|
||||
}
|
||||
};
|
||||
|
||||
let envelope = RoomMessageEnvelope {
|
||||
id: streaming_msg_id,
|
||||
dedup_key: Some(format!("{}:{}", room_id_inner, streaming_msg_id)),
|
||||
room_id: room_id_inner,
|
||||
sender_type: sender_type.clone(),
|
||||
sender_id: None,
|
||||
model_id: Some(model_id),
|
||||
thread_id: None,
|
||||
content: persist_content.clone(),
|
||||
content_type: "text".to_string(),
|
||||
thinking_content,
|
||||
send_at: now,
|
||||
seq,
|
||||
in_reply_to: None,
|
||||
display_name: Some(ai_display_name.clone()),
|
||||
};
|
||||
|
||||
if let Err(e) = queue.publish(room_id_inner, envelope).await {
|
||||
tracing::error!(error = %e, "Failed to publish ReAct streaming message");
|
||||
} else {
|
||||
let now = Utc::now();
|
||||
if let Err(e) = room_ai::Entity::update_many()
|
||||
.col_expr(
|
||||
room_ai::Column::CallCount,
|
||||
Expr::col(room_ai::Column::CallCount).add(1),
|
||||
)
|
||||
.col_expr(room_ai::Column::LastCallAt, Expr::value(Some(now)))
|
||||
.filter(room_ai::Column::Room.eq(room_id_inner))
|
||||
.filter(room_ai::Column::Model.eq(model_id))
|
||||
.exec(&db)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "Failed to update room_ai call stats");
|
||||
}
|
||||
|
||||
// Record billing (non-fatal)
|
||||
// TODO: ReAct agent does not track token counts yet; billing with 0/0
|
||||
let _ = super::billing::record_ai_usage(
|
||||
&db,
|
||||
project_id_inner,
|
||||
model_id,
|
||||
0,
|
||||
0,
|
||||
)
|
||||
.await;
|
||||
|
||||
let msg_event = queue::RoomMessageEvent {
|
||||
id: streaming_msg_id,
|
||||
room_id: room_id_inner,
|
||||
sender_type: sender_type.clone(),
|
||||
sender_id: None,
|
||||
thread_id: None,
|
||||
content: persist_content,
|
||||
content_type: "text".to_string(),
|
||||
thinking_content: None,
|
||||
send_at: now,
|
||||
seq,
|
||||
display_name: Some(ai_display_name.clone()),
|
||||
in_reply_to: None,
|
||||
reactions: None,
|
||||
message_id: None,
|
||||
};
|
||||
room_manager.broadcast(room_id_inner, msg_event).await;
|
||||
room_manager.metrics.messages_sent.increment(1);
|
||||
|
||||
let event = ProjectRoomEvent {
|
||||
event_type: crate::RoomEventType::NewMessage.as_str().into(),
|
||||
project_id: project_id_inner,
|
||||
room_id: Some(room_id_inner),
|
||||
category_id: None,
|
||||
message_id: Some(streaming_msg_id),
|
||||
seq: Some(seq),
|
||||
timestamp: now,
|
||||
};
|
||||
queue
|
||||
.publish_project_room_event(project_id_inner, event)
|
||||
.await;
|
||||
}
|
||||
|
||||
room_manager.close_stream_channel(streaming_msg_id).await;
|
||||
});
|
||||
}
|
||||
274
libs/room/src/service/ai_streaming.rs
Normal file
274
libs/room/src/service/ai_streaming.rs
Normal file
@ -0,0 +1,274 @@
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room_ai;
|
||||
use queue::{MessageProducer, ProjectRoomEvent, RoomMessageEnvelope};
|
||||
use sea_orm::{sea_query::Expr, ColumnTrait, EntityTrait, ExprTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::sequence::next_room_message_seq_internal;
|
||||
use crate::connection::RoomConnectionManager;
|
||||
use agent::chat::{AiRequest, ChatService};
|
||||
|
||||
pub async fn process_message_ai_streaming(
|
||||
chat_service: Arc<ChatService>,
|
||||
request: AiRequest,
|
||||
room_id: Uuid,
|
||||
project_id: Uuid,
|
||||
model_id: Uuid,
|
||||
lock_guard: crate::room_ai_queue::RoomAiLockGuard,
|
||||
db: AppDatabase,
|
||||
cache: AppCache,
|
||||
queue: MessageProducer,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
) {
|
||||
use queue::RoomMessageStreamChunkEvent;
|
||||
|
||||
let streaming_msg_id = Uuid::now_v7();
|
||||
let seq = match next_room_message_seq_internal(room_id, &db, &cache).await {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "Failed to get seq for streaming AI message");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let _ = room_manager
|
||||
.register_stream_channel(streaming_msg_id)
|
||||
.await;
|
||||
|
||||
let initial_event = RoomMessageStreamChunkEvent {
|
||||
message_id: streaming_msg_id,
|
||||
room_id,
|
||||
content: String::new(),
|
||||
done: false,
|
||||
error: None,
|
||||
display_name: Some(request.model.name.clone()),
|
||||
chunk_type: Some("thinking".to_string()),
|
||||
};
|
||||
room_manager.broadcast_stream_chunk(initial_event).await;
|
||||
|
||||
let room_id_inner = room_id;
|
||||
let project_id_inner = project_id;
|
||||
let now = Utc::now();
|
||||
let sender_type = "ai".to_string();
|
||||
let ai_display_name = request.model.name.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _lock_guard = lock_guard;
|
||||
let ai_typing_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
|
||||
let ai_display_name_for_chunk = ai_display_name.clone();
|
||||
let ai_display_name_for_final = ai_display_name.clone();
|
||||
|
||||
let chunk_count = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
|
||||
let room_manager_cb = room_manager.clone();
|
||||
|
||||
let on_chunk = move |chunk: agent::chat::AiStreamChunk| {
|
||||
Box::pin({
|
||||
let room_manager = room_manager_cb.clone();
|
||||
let streaming_msg_id = streaming_msg_id;
|
||||
let room_id = room_id_inner;
|
||||
let chunk_count = chunk_count.clone();
|
||||
let ai_display_name_for_chunk = ai_display_name_for_chunk.clone();
|
||||
async move {
|
||||
let chunk_type_str = match chunk.chunk_type {
|
||||
agent::chat::AiChunkType::Thinking => "thinking",
|
||||
agent::chat::AiChunkType::Answer => "answer",
|
||||
agent::chat::AiChunkType::ToolCall => "tool_call",
|
||||
agent::chat::AiChunkType::ToolResult => "tool_result",
|
||||
};
|
||||
let event = RoomMessageStreamChunkEvent {
|
||||
message_id: streaming_msg_id,
|
||||
room_id,
|
||||
content: chunk.content,
|
||||
done: chunk.done,
|
||||
error: None,
|
||||
display_name: Some(ai_display_name_for_chunk),
|
||||
chunk_type: Some(chunk_type_str.to_string()),
|
||||
};
|
||||
room_manager.broadcast_stream_chunk(event).await;
|
||||
chunk_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
}) as Pin<Box<dyn std::future::Future<Output = ()> + Send>>
|
||||
};
|
||||
|
||||
let stream_callback: agent::chat::StreamCallback = Box::new(on_chunk);
|
||||
|
||||
let typing_start = queue::TypingEvent {
|
||||
room_id: room_id_inner,
|
||||
user_id: ai_typing_id,
|
||||
username: ai_display_name.clone(),
|
||||
avatar_url: None,
|
||||
action: "start".to_string(),
|
||||
sender_type: Some("ai".to_string()),
|
||||
};
|
||||
room_manager.broadcast_typing(room_id_inner, typing_start.clone()).await;
|
||||
|
||||
let (typing_cancel_tx, typing_cancel_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
let typing_renew_handle = tokio::spawn({
|
||||
let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
let mgr = room_manager.clone();
|
||||
let rid = room_id_inner;
|
||||
let evt = typing_start.clone();
|
||||
async move {
|
||||
tokio::select! {
|
||||
_ = typing_cancel_rx => {}
|
||||
_ = async {
|
||||
loop {
|
||||
interval.tick().await;
|
||||
mgr.broadcast_typing(rid, evt.clone()).await;
|
||||
}
|
||||
} => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
match chat_service.process_stream(request, stream_callback).await {
|
||||
Ok(result) => {
|
||||
// Store ordered chunks as JSON in thinking_content for ordered replay.
|
||||
// Uses {"__chunks__": [...]} marker so legacy plain-text still works.
|
||||
let thinking_content = if result.chunks.is_empty() {
|
||||
None
|
||||
} else {
|
||||
let chunks_json = serde_json::json!({
|
||||
"__chunks__": result.chunks.iter().map(|c| {
|
||||
let type_str = match c.chunk_type {
|
||||
agent::client::StreamChunkType::Thinking => "thinking",
|
||||
agent::client::StreamChunkType::Answer => "answer",
|
||||
agent::client::StreamChunkType::ToolCall => "tool_call",
|
||||
};
|
||||
serde_json::json!({
|
||||
"type": type_str,
|
||||
"content": c.content,
|
||||
})
|
||||
}).collect::<Vec<_>>(),
|
||||
});
|
||||
Some(chunks_json.to_string())
|
||||
};
|
||||
let envelope = RoomMessageEnvelope {
|
||||
id: streaming_msg_id,
|
||||
dedup_key: Some(format!("{}:{}", room_id_inner, streaming_msg_id)),
|
||||
room_id: room_id_inner,
|
||||
sender_type: sender_type.clone(),
|
||||
sender_id: None,
|
||||
model_id: Some(model_id),
|
||||
thread_id: None,
|
||||
content: result.content.clone(),
|
||||
content_type: "text".to_string(),
|
||||
thinking_content: thinking_content.clone(),
|
||||
send_at: now,
|
||||
seq,
|
||||
in_reply_to: None,
|
||||
display_name: Some(ai_display_name_for_final.clone()),
|
||||
};
|
||||
|
||||
if let Err(e) = queue.publish(room_id_inner, envelope).await {
|
||||
tracing::error!(error = %e, "Failed to publish streaming AI message");
|
||||
} else {
|
||||
let now = Utc::now();
|
||||
if let Err(e) = room_ai::Entity::update_many()
|
||||
.col_expr(
|
||||
room_ai::Column::CallCount,
|
||||
Expr::col(room_ai::Column::CallCount).add(1),
|
||||
)
|
||||
.col_expr(
|
||||
room_ai::Column::LastCallAt,
|
||||
Expr::value(Some(now)),
|
||||
)
|
||||
.filter(room_ai::Column::Room.eq(room_id_inner))
|
||||
.filter(room_ai::Column::Model.eq(model_id))
|
||||
.exec(&db)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "Failed to update room_ai call stats");
|
||||
}
|
||||
|
||||
// Record billing (non-fatal)
|
||||
let _ = super::billing::record_ai_usage(
|
||||
&db,
|
||||
project_id_inner,
|
||||
model_id,
|
||||
result.input_tokens,
|
||||
result.output_tokens,
|
||||
)
|
||||
.await;
|
||||
|
||||
let msg_event = queue::RoomMessageEvent {
|
||||
id: streaming_msg_id,
|
||||
room_id: room_id_inner,
|
||||
sender_type: sender_type.clone(),
|
||||
sender_id: None,
|
||||
thread_id: None,
|
||||
content: result.content.clone(),
|
||||
content_type: "text".to_string(),
|
||||
thinking_content: thinking_content.clone(),
|
||||
send_at: now,
|
||||
seq,
|
||||
display_name: Some(ai_display_name_for_final.clone()),
|
||||
in_reply_to: None,
|
||||
reactions: None,
|
||||
message_id: None,
|
||||
};
|
||||
room_manager.broadcast(room_id_inner, msg_event).await;
|
||||
room_manager.metrics.messages_sent.increment(1);
|
||||
|
||||
let _ = typing_cancel_tx.send(());
|
||||
typing_renew_handle.abort();
|
||||
let typing_stop = queue::TypingEvent {
|
||||
room_id: room_id_inner,
|
||||
user_id: ai_typing_id,
|
||||
username: ai_display_name_for_final.clone(),
|
||||
avatar_url: None,
|
||||
action: "stop".to_string(),
|
||||
sender_type: Some("ai".to_string()),
|
||||
};
|
||||
room_manager.broadcast_typing(room_id_inner, typing_stop).await;
|
||||
|
||||
let event = ProjectRoomEvent {
|
||||
event_type: crate::RoomEventType::NewMessage.as_str().into(),
|
||||
project_id: project_id_inner,
|
||||
room_id: Some(room_id_inner),
|
||||
category_id: None,
|
||||
message_id: Some(streaming_msg_id),
|
||||
seq: Some(seq),
|
||||
timestamp: now,
|
||||
};
|
||||
queue
|
||||
.publish_project_room_event(project_id_inner, event)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "AI streaming failed");
|
||||
let _ = typing_cancel_tx.send(());
|
||||
typing_renew_handle.abort();
|
||||
let typing_stop = queue::TypingEvent {
|
||||
room_id: room_id_inner,
|
||||
user_id: ai_typing_id,
|
||||
username: ai_display_name.clone(),
|
||||
avatar_url: None,
|
||||
action: "stop".to_string(),
|
||||
sender_type: Some("ai".to_string()),
|
||||
};
|
||||
room_manager.broadcast_typing(room_id_inner, typing_stop).await;
|
||||
|
||||
let event = RoomMessageStreamChunkEvent {
|
||||
message_id: streaming_msg_id,
|
||||
room_id: room_id_inner,
|
||||
content: String::new(),
|
||||
done: true,
|
||||
error: Some(e.to_string()),
|
||||
display_name: Some(ai_display_name.clone()),
|
||||
chunk_type: None,
|
||||
};
|
||||
room_manager.broadcast_stream_chunk(event).await;
|
||||
}
|
||||
}
|
||||
|
||||
room_manager.close_stream_channel(streaming_msg_id).await;
|
||||
});
|
||||
}
|
||||
51
libs/room/src/service/billing.rs
Normal file
51
libs/room/src/service/billing.rs
Normal file
@ -0,0 +1,51 @@
|
||||
//! AI usage billing helper for room service.
|
||||
//!
|
||||
//! Delegates to `agent::billing::record_ai_usage`.
|
||||
//! Billing is non-fatal — failures are logged but do not block AI responses.
|
||||
|
||||
use db::database::AppDatabase;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::RoomError;
|
||||
|
||||
/// Record AI token usage against a project's billing balance.
|
||||
///
|
||||
/// Returns `Ok(())` on success. On billing failure (e.g. insufficient balance,
|
||||
/// missing pricing), returns `Err` but the caller should still complete the AI
|
||||
/// request — billing is a non-critical side-effect.
|
||||
pub async fn record_ai_usage(
|
||||
db: &AppDatabase,
|
||||
project_id: Uuid,
|
||||
model_id: Uuid,
|
||||
input_tokens: i64,
|
||||
output_tokens: i64,
|
||||
) -> Result<(), RoomError> {
|
||||
if input_tokens == 0 && output_tokens == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
match agent::billing::record_ai_usage(db, project_id, model_id, input_tokens, output_tokens).await {
|
||||
Ok(record) => {
|
||||
tracing::info!(
|
||||
project_id = %project_id,
|
||||
model_id = %model_id,
|
||||
input_tokens = input_tokens,
|
||||
output_tokens = output_tokens,
|
||||
cost_usd = %record.cost,
|
||||
"ai_usage_recorded"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
project_id = %project_id,
|
||||
model_id = %model_id,
|
||||
input_tokens = input_tokens,
|
||||
output_tokens = output_tokens,
|
||||
error = %e,
|
||||
"ai_billing_failed_non_fatal"
|
||||
);
|
||||
Err(e.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
62
libs/room/src/service/history.rs
Normal file
62
libs/room/src/service/history.rs
Normal file
@ -0,0 +1,62 @@
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room_ai;
|
||||
use models::rooms::room_message::{Column as RmCol, Entity as RoomMessage};
|
||||
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::RoomError;
|
||||
|
||||
pub async fn get_room_history(
|
||||
db: &AppDatabase,
|
||||
room_id: Uuid,
|
||||
limit: usize,
|
||||
) -> Result<Vec<models::rooms::room_message::Model>, RoomError> {
|
||||
let messages = RoomMessage::find()
|
||||
.filter(RmCol::Room.eq(room_id))
|
||||
.order_by_desc(RmCol::Seq)
|
||||
.limit(limit as u64)
|
||||
.all(db)
|
||||
.await?;
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
|
||||
pub async fn get_user_names(
|
||||
db: &AppDatabase,
|
||||
user_ids: &[Uuid],
|
||||
) -> std::collections::HashMap<Uuid, String> {
|
||||
use models::users::User;
|
||||
|
||||
let mut names = std::collections::HashMap::new();
|
||||
if user_ids.is_empty() {
|
||||
return names;
|
||||
}
|
||||
|
||||
let users = User::find()
|
||||
.filter(models::users::user::Column::Uid.is_in(user_ids.to_vec()))
|
||||
.all(db)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
for user in users {
|
||||
names.insert(user.uid, user.username);
|
||||
}
|
||||
|
||||
names
|
||||
}
|
||||
|
||||
pub async fn get_room_ai_config(
|
||||
db: &AppDatabase,
|
||||
room_id: Uuid,
|
||||
) -> Result<Option<room_ai::Model>, RoomError> {
|
||||
let ai_config = room_ai::Entity::find()
|
||||
.filter(room_ai::Column::Room.eq(room_id))
|
||||
.one(db)
|
||||
.await?;
|
||||
|
||||
Ok(ai_config)
|
||||
}
|
||||
|
||||
pub async fn extract_mention_context(_content: &str) -> Vec<agent::chat::Mention> {
|
||||
Vec::new()
|
||||
}
|
||||
48
libs/room/src/service/mentions.rs
Normal file
48
libs/room/src/service/mentions.rs
Normal file
@ -0,0 +1,48 @@
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::patterns::{mention_bracket_re, mention_tag_re, user_mention_re};
|
||||
|
||||
/// Extracts user UUIDs from all mention formats:
|
||||
/// - Legacy: `<user>uuid</user>`
|
||||
/// - Legacy: `<mention type="user" id="uuid">label</mention>`
|
||||
/// - New: `@[user:uuid:label]`
|
||||
pub fn extract_mentions(content: &str) -> Vec<Uuid> {
|
||||
let mut mentioned = Vec::new();
|
||||
|
||||
for cap in user_mention_re().captures_iter(content) {
|
||||
if let Some(inner) = cap.get(1) {
|
||||
let token = inner.as_str().trim();
|
||||
if let Ok(uuid) = Uuid::parse_str(token) {
|
||||
if !mentioned.contains(&uuid) {
|
||||
mentioned.push(uuid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for cap in mention_tag_re().captures_iter(content) {
|
||||
if let (Some(type_m), Some(id_m)) = (cap.get(1), cap.get(2)) {
|
||||
if type_m.as_str() == "user" {
|
||||
if let Ok(uuid) = Uuid::parse_str(id_m.as_str().trim()) {
|
||||
if !mentioned.contains(&uuid) {
|
||||
mentioned.push(uuid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for cap in mention_bracket_re().captures_iter(content) {
|
||||
if let (Some(type_m), Some(id_m)) = (cap.get(1), cap.get(2)) {
|
||||
if type_m.as_str() == "user" {
|
||||
if let Ok(uuid) = Uuid::parse_str(id_m.as_str().trim()) {
|
||||
if !mentioned.contains(&uuid) {
|
||||
mentioned.push(uuid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
mentioned
|
||||
}
|
||||
466
libs/room/src/service/mod.rs
Normal file
466
libs/room/src/service/mod.rs
Normal file
@ -0,0 +1,466 @@
|
||||
mod access;
|
||||
mod billing;
|
||||
mod ai_common;
|
||||
mod ai_nonstreaming;
|
||||
mod ai_react_nonstreaming;
|
||||
mod ai_react_streaming;
|
||||
mod ai_streaming;
|
||||
mod history;
|
||||
mod mentions;
|
||||
mod notifications;
|
||||
mod patterns;
|
||||
mod sequence;
|
||||
mod workers;
|
||||
|
||||
pub use access::{check_room_access, check_project_member, require_room_member, find_room_or_404};
|
||||
pub use ai_common::create_and_publish_ai_message;
|
||||
pub use ai_nonstreaming::process_message_ai_nonstreaming;
|
||||
pub use ai_react_nonstreaming::process_message_ai_react_nonstreaming;
|
||||
pub use ai_react_streaming::process_message_ai_react_streaming;
|
||||
pub use ai_streaming::process_message_ai_streaming;
|
||||
pub use history::{get_room_history, get_user_names, get_room_ai_config, extract_mention_context};
|
||||
pub use mentions::extract_mentions;
|
||||
pub use notifications::{notify_project_members, publish_room_event};
|
||||
pub use sequence::next_room_message_seq_internal;
|
||||
pub use workers::{start_workers, spawn_agent_task, spawn_room_workers, PushNotificationFn};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room;
|
||||
use models::rooms::room_ai;
|
||||
use queue::{MessageProducer, ProjectRoomEvent};
|
||||
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::connection::{RoomConnectionManager, DedupCache};
|
||||
use crate::error::RoomError;
|
||||
use agent::chat::{AiRequest, ChatService};
|
||||
use agent::embed::EmbedService;
|
||||
use agent::TaskService;
|
||||
use models::agent_task::AgentType;
|
||||
|
||||
use crate::service::patterns::{mention_bracket_re, mention_tag_re};
|
||||
|
||||
const DEFAULT_MAX_CONCURRENT_WORKERS: usize = 1024;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RoomService {
|
||||
pub db: AppDatabase,
|
||||
pub cache: AppCache,
|
||||
pub config: config::AppConfig,
|
||||
pub room_manager: Arc<RoomConnectionManager>,
|
||||
pub queue: MessageProducer,
|
||||
pub redis_url: String,
|
||||
pub chat_service: Option<Arc<ChatService>>,
|
||||
pub task_service: Option<Arc<TaskService>>,
|
||||
pub embed_service: Option<Arc<EmbedService>>,
|
||||
pub push_fn: Option<workers::PushNotificationFn>,
|
||||
worker_semaphore: Arc<tokio::sync::Semaphore>,
|
||||
dedup_cache: DedupCache,
|
||||
}
|
||||
|
||||
impl RoomService {
|
||||
pub fn new(
|
||||
db: AppDatabase,
|
||||
cache: AppCache,
|
||||
config: config::AppConfig,
|
||||
queue: MessageProducer,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
redis_url: String,
|
||||
chat_service: Option<Arc<ChatService>>,
|
||||
task_service: Option<Arc<TaskService>>,
|
||||
max_concurrent_workers: Option<usize>,
|
||||
push_fn: Option<workers::PushNotificationFn>,
|
||||
embed_service: Option<Arc<EmbedService>>,
|
||||
) -> Self {
|
||||
let dedup_cache: DedupCache =
|
||||
Arc::new(dashmap::DashMap::with_capacity_and_hasher(10000, Default::default()));
|
||||
Self {
|
||||
db,
|
||||
cache,
|
||||
config,
|
||||
room_manager,
|
||||
queue,
|
||||
redis_url,
|
||||
chat_service,
|
||||
task_service,
|
||||
embed_service,
|
||||
worker_semaphore: Arc::new(tokio::sync::Semaphore::new(
|
||||
max_concurrent_workers.unwrap_or(DEFAULT_MAX_CONCURRENT_WORKERS),
|
||||
)),
|
||||
dedup_cache,
|
||||
push_fn,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn start_workers(
|
||||
&self,
|
||||
shutdown_rx: tokio::sync::broadcast::Receiver<()>,
|
||||
) -> anyhow::Result<()> {
|
||||
workers::start_workers(
|
||||
self.db.clone(),
|
||||
self.cache.clone(),
|
||||
self.room_manager.clone(),
|
||||
self.queue.clone(),
|
||||
self.redis_url.clone(),
|
||||
self.dedup_cache.clone(),
|
||||
self.task_service.clone(),
|
||||
None, // max_concurrent_workers handled by semaphore
|
||||
shutdown_rx,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn spawn_agent_task<F, Fut>(
|
||||
&self,
|
||||
project_id: Uuid,
|
||||
agent_type: AgentType,
|
||||
input: String,
|
||||
_title: Option<String>,
|
||||
execute: F,
|
||||
) -> anyhow::Result<i64>
|
||||
where
|
||||
F: FnOnce(i64, Arc<TaskService>) -> Fut + Send + 'static,
|
||||
Fut: std::future::Future<Output = Result<String, String>> + Send,
|
||||
{
|
||||
let task_service = match &self.task_service {
|
||||
Some(ts) => ts.clone(),
|
||||
None => return Err(anyhow::anyhow!("task service not configured")),
|
||||
};
|
||||
|
||||
workers::spawn_agent_task(
|
||||
project_id,
|
||||
agent_type,
|
||||
input,
|
||||
task_service,
|
||||
self.queue.clone(),
|
||||
self.room_manager.clone(),
|
||||
self.worker_semaphore.clone(),
|
||||
execute,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn spawn_room_workers(&self, room_id: uuid::Uuid) {
|
||||
workers::spawn_room_workers(
|
||||
room_id,
|
||||
self.db.clone(),
|
||||
self.room_manager.clone(),
|
||||
self.queue.clone(),
|
||||
self.redis_url.clone(),
|
||||
self.worker_semaphore.clone(),
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn publish_room_event(
|
||||
&self,
|
||||
project_id: uuid::Uuid,
|
||||
event_type: super::RoomEventType,
|
||||
room_id: Option<uuid::Uuid>,
|
||||
category_id: Option<uuid::Uuid>,
|
||||
message_id: Option<uuid::Uuid>,
|
||||
seq: Option<i64>,
|
||||
) {
|
||||
let event = ProjectRoomEvent {
|
||||
event_type: event_type.as_str().into(),
|
||||
project_id,
|
||||
room_id,
|
||||
category_id,
|
||||
message_id,
|
||||
seq,
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
self.queue
|
||||
.publish_project_room_event(project_id, event)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub fn notify_project_members(
|
||||
&self,
|
||||
project_id: uuid::Uuid,
|
||||
notification_type: super::NotificationType,
|
||||
title: String,
|
||||
content: Option<String>,
|
||||
related_room_id: Option<uuid::Uuid>,
|
||||
) {
|
||||
notifications::notify_project_members(
|
||||
self.db.clone(),
|
||||
project_id,
|
||||
notification_type,
|
||||
title,
|
||||
content,
|
||||
related_room_id,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn extract_mentions(content: &str) -> Vec<Uuid> {
|
||||
mentions::extract_mentions(content)
|
||||
}
|
||||
|
||||
pub async fn resolve_mentions(&self, content: &str) -> Vec<Uuid> {
|
||||
use models::users::User;
|
||||
use sea_orm::EntityTrait;
|
||||
|
||||
let mut resolved: Vec<Uuid> = Vec::new();
|
||||
let mut seen_usernames: Vec<String> = Vec::new();
|
||||
|
||||
for cap in mention_bracket_re().captures_iter(content) {
|
||||
if let (Some(type_m), Some(id_m)) = (cap.get(1), cap.get(2)) {
|
||||
if type_m.as_str() == "user" {
|
||||
let id = id_m.as_str().trim();
|
||||
if let Ok(uuid) = Uuid::parse_str(id) {
|
||||
if !resolved.contains(&uuid) {
|
||||
resolved.push(uuid);
|
||||
}
|
||||
} else if let Some(label_m) = cap.get(3) {
|
||||
let label = label_m.as_str().trim();
|
||||
if !label.is_empty() {
|
||||
let label_lower = label.to_lowercase();
|
||||
if seen_usernames.contains(&label_lower) {
|
||||
continue;
|
||||
}
|
||||
seen_usernames.push(label_lower.clone());
|
||||
|
||||
if let Some(user) = User::find()
|
||||
.filter(models::users::user::Column::Username.eq(label_lower))
|
||||
.one(&self.db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
{
|
||||
if !resolved.contains(&user.uid) {
|
||||
resolved.push(user.uid);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
resolved
|
||||
}
|
||||
|
||||
pub async fn check_room_access(&self, room_id: Uuid, user_id: Uuid) -> Result<(), RoomError> {
|
||||
access::check_room_access(&self.db, room_id, user_id).await
|
||||
}
|
||||
|
||||
pub async fn check_project_member(
|
||||
&self,
|
||||
project_id: Uuid,
|
||||
user_id: Uuid,
|
||||
) -> Result<(), RoomError> {
|
||||
access::check_project_member(&self.db, project_id, user_id).await
|
||||
}
|
||||
|
||||
pub async fn should_ai_respond(&self, room_id: Uuid, content: &str) -> Result<bool, RoomError> {
|
||||
let ai_config = history::get_room_ai_config(&self.db, room_id).await?;
|
||||
|
||||
let config = match ai_config {
|
||||
Some(c) => c,
|
||||
None => return Ok(false),
|
||||
};
|
||||
|
||||
if !config.use_exact {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let model_id_str = config.model.to_string();
|
||||
|
||||
for cap in mention_bracket_re().captures_iter(content) {
|
||||
if let (Some(type_m), Some(id_m)) = (cap.get(1), cap.get(2)) {
|
||||
if type_m.as_str() == "ai" && id_m.as_str().trim() == model_id_str {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for cap in mention_tag_re().captures_iter(content) {
|
||||
if let (Some(type_m), Some(id_m)) = (cap.get(1), cap.get(2)) {
|
||||
if type_m.as_str() == "ai" && id_m.as_str().trim() == model_id_str {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
pub async fn get_room_ai_config(
|
||||
&self,
|
||||
room_id: Uuid,
|
||||
) -> Result<Option<room_ai::Model>, RoomError> {
|
||||
history::get_room_ai_config(&self.db, room_id).await
|
||||
}
|
||||
|
||||
pub async fn get_user_names(
|
||||
&self,
|
||||
user_ids: &[Uuid],
|
||||
) -> std::collections::HashMap<Uuid, String> {
|
||||
history::get_user_names(&self.db, user_ids).await
|
||||
}
|
||||
|
||||
pub async fn require_room_member(&self, room_id: Uuid, user_id: Uuid) -> Result<(), RoomError> {
|
||||
access::require_room_member(&self.db, room_id, user_id).await
|
||||
}
|
||||
|
||||
pub async fn find_room_or_404(&self, room_id: Uuid) -> Result<room::Model, RoomError> {
|
||||
access::find_room_or_404(&self.db, room_id).await
|
||||
}
|
||||
|
||||
pub async fn process_message_ai(
|
||||
&self,
|
||||
room_id: Uuid,
|
||||
_message_id: Uuid,
|
||||
sender_id: Uuid,
|
||||
content: String,
|
||||
) -> Result<(), RoomError> {
|
||||
let Some(chat_service) = &self.chat_service else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let Some(ai_config) = self.get_room_ai_config(room_id).await? else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let Some(lock_guard) =
|
||||
crate::room_ai_queue::acquire_room_ai_lock(&self.cache, room_id).await?
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let room = self.find_room_or_404(room_id).await?;
|
||||
|
||||
let project = models::projects::project::Entity::find_by_id(room.project)
|
||||
.one(&self.db)
|
||||
.await?
|
||||
.ok_or_else(|| RoomError::NotFound("Project not found".to_string()))?;
|
||||
|
||||
let mentioned_model_id = {
|
||||
let mut found = None;
|
||||
for cap in mention_bracket_re().captures_iter(&content) {
|
||||
if let (Some(type_m), Some(id_m)) = (cap.get(1), cap.get(2)) {
|
||||
if type_m.as_str() == "ai" {
|
||||
if let Ok(uuid) = Uuid::parse_str(id_m.as_str().trim()) {
|
||||
found = Some(uuid);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
found
|
||||
};
|
||||
|
||||
let model_id = mentioned_model_id.unwrap_or(ai_config.model);
|
||||
let model = models::agents::model::Entity::find_by_id(model_id)
|
||||
.one(&self.db)
|
||||
.await?
|
||||
.ok_or_else(|| RoomError::NotFound("AI model not found".to_string()))?;
|
||||
|
||||
let sender = models::users::User::find_by_id(sender_id)
|
||||
.one(&self.db)
|
||||
.await?
|
||||
.ok_or_else(|| RoomError::NotFound("Sender not found".to_string()))?;
|
||||
|
||||
let history = history::get_room_history(&self.db, room_id, 50).await?;
|
||||
|
||||
let user_ids: Vec<Uuid> = history
|
||||
.iter()
|
||||
.filter_map(|m| m.sender_id)
|
||||
.chain(std::iter::once(sender_id))
|
||||
.collect();
|
||||
let user_names = self.get_user_names(&user_ids).await;
|
||||
|
||||
let mentions = history::extract_mention_context(&content).await;
|
||||
|
||||
let request = AiRequest {
|
||||
db: self.db.clone(),
|
||||
cache: self.cache.clone(),
|
||||
config: self.config.clone(),
|
||||
model,
|
||||
project: project.clone(),
|
||||
sender,
|
||||
room: room.clone(),
|
||||
input: content,
|
||||
mention: mentions,
|
||||
history,
|
||||
user_names,
|
||||
temperature: ai_config.temperature.unwrap_or(0.7),
|
||||
max_tokens: ai_config.max_tokens.unwrap_or(4096) as i32,
|
||||
top_p: 1.0,
|
||||
frequency_penalty: 0.0,
|
||||
presence_penalty: 0.0,
|
||||
think: ai_config.think,
|
||||
tools: Some(chat_service.tools()),
|
||||
max_tool_depth: 1000,
|
||||
};
|
||||
|
||||
let use_streaming = ai_config.stream;
|
||||
let is_react = ai_config.agent_type.as_deref() == Some("react");
|
||||
|
||||
if is_react {
|
||||
if use_streaming {
|
||||
ai_react_streaming::process_message_ai_react_streaming(
|
||||
chat_service.clone(),
|
||||
request,
|
||||
room_id,
|
||||
room.project,
|
||||
model_id,
|
||||
lock_guard,
|
||||
self.db.clone(),
|
||||
self.cache.clone(),
|
||||
self.queue.clone(),
|
||||
self.room_manager.clone(),
|
||||
)
|
||||
.await;
|
||||
} else {
|
||||
ai_react_nonstreaming::process_message_ai_react_nonstreaming(
|
||||
chat_service.clone(),
|
||||
request,
|
||||
room_id,
|
||||
room.project,
|
||||
model_id,
|
||||
lock_guard,
|
||||
self.db.clone(),
|
||||
self.cache.clone(),
|
||||
self.queue.clone(),
|
||||
self.room_manager.clone(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
} else if use_streaming {
|
||||
ai_streaming::process_message_ai_streaming(
|
||||
chat_service.clone(),
|
||||
request,
|
||||
room_id,
|
||||
room.project,
|
||||
model_id,
|
||||
lock_guard,
|
||||
self.db.clone(),
|
||||
self.cache.clone(),
|
||||
self.queue.clone(),
|
||||
self.room_manager.clone(),
|
||||
)
|
||||
.await;
|
||||
} else {
|
||||
ai_nonstreaming::process_message_ai_nonstreaming(
|
||||
chat_service.clone(),
|
||||
request,
|
||||
room_id,
|
||||
room.project,
|
||||
model_id,
|
||||
lock_guard,
|
||||
self.db.clone(),
|
||||
self.cache.clone(),
|
||||
self.queue.clone(),
|
||||
self.room_manager.clone(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
134
libs/room/src/service/notifications.rs
Normal file
134
libs/room/src/service/notifications.rs
Normal file
@ -0,0 +1,134 @@
|
||||
use chrono::Utc;
|
||||
use db::database::AppDatabase;
|
||||
use models::projects::project_members;
|
||||
use queue::ProjectRoomEvent;
|
||||
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::RoomError;
|
||||
|
||||
pub fn notify_project_members(
|
||||
db: AppDatabase,
|
||||
project_id: Uuid,
|
||||
notification_type: crate::NotificationType,
|
||||
title: String,
|
||||
content: Option<String>,
|
||||
related_room_id: Option<Uuid>,
|
||||
) {
|
||||
let notification_type_inner = notification_type;
|
||||
let title_inner = title;
|
||||
let content_inner = content;
|
||||
let related_room_id_inner = related_room_id;
|
||||
let project_id_inner = project_id;
|
||||
|
||||
tokio::spawn(async move {
|
||||
let members = match project_members::Entity::find()
|
||||
.filter(project_members::Column::Project.eq(project_id_inner))
|
||||
.all(&db)
|
||||
.await
|
||||
{
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
tracing::error!(project_id = %project_id_inner, error = %e,
|
||||
"notify_project_members: failed to fetch members");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
for member in members {
|
||||
let user_id = member.user;
|
||||
if let Err(e) = create_notification_sync(
|
||||
&db,
|
||||
notification_type_inner,
|
||||
user_id,
|
||||
title_inner.clone(),
|
||||
content_inner.clone(),
|
||||
related_room_id_inner,
|
||||
project_id_inner,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(user_id = %user_id, project_id = %project_id_inner, error = %e,
|
||||
"notify_project_members: failed to create notification for user");
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async fn create_notification_sync(
|
||||
db: &AppDatabase,
|
||||
notification_type: crate::NotificationType,
|
||||
user_id: Uuid,
|
||||
title: String,
|
||||
content: Option<String>,
|
||||
related_room_id: Option<Uuid>,
|
||||
project_id: Uuid,
|
||||
) -> Result<(), RoomError> {
|
||||
use models::rooms::room_notifications;
|
||||
use sea_orm::{ActiveModelTrait, Set};
|
||||
|
||||
let notification_type_model = match notification_type {
|
||||
crate::NotificationType::Mention => room_notifications::NotificationType::Mention,
|
||||
crate::NotificationType::Invitation => room_notifications::NotificationType::Invitation,
|
||||
crate::NotificationType::RoleChange => room_notifications::NotificationType::RoleChange,
|
||||
crate::NotificationType::RoomCreated => room_notifications::NotificationType::RoomCreated,
|
||||
crate::NotificationType::RoomDeleted => room_notifications::NotificationType::RoomDeleted,
|
||||
crate::NotificationType::SystemAnnouncement => {
|
||||
room_notifications::NotificationType::SystemAnnouncement
|
||||
}
|
||||
crate::NotificationType::ProjectInvitation => {
|
||||
room_notifications::NotificationType::ProjectInvitation
|
||||
}
|
||||
crate::NotificationType::WorkspaceInvitation => {
|
||||
room_notifications::NotificationType::WorkspaceInvitation
|
||||
}
|
||||
};
|
||||
|
||||
let _model = room_notifications::ActiveModel {
|
||||
id: Set(Uuid::now_v7()),
|
||||
room: Set(related_room_id),
|
||||
project: Set(Some(project_id)),
|
||||
user_id: Set(Some(user_id)),
|
||||
notification_type: Set(notification_type_model),
|
||||
related_message_id: Set(None),
|
||||
related_user_id: Set(None),
|
||||
related_room_id: Set(related_room_id),
|
||||
title: Set(title),
|
||||
content: Set(content),
|
||||
metadata: Set(None),
|
||||
is_read: Set(false),
|
||||
is_archived: Set(false),
|
||||
created_at: Set(Utc::now()),
|
||||
read_at: Set(None),
|
||||
expires_at: Set(None),
|
||||
}
|
||||
.insert(db)
|
||||
.await
|
||||
.map_err(|e| RoomError::Database(e))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn publish_room_event(
|
||||
queue: &queue::MessageProducer,
|
||||
project_id: Uuid,
|
||||
event_type: crate::RoomEventType,
|
||||
room_id: Option<Uuid>,
|
||||
message_id: Option<Uuid>,
|
||||
seq: Option<i64>,
|
||||
) {
|
||||
let event = ProjectRoomEvent {
|
||||
event_type: event_type.as_str().into(),
|
||||
project_id,
|
||||
room_id,
|
||||
category_id: None,
|
||||
message_id,
|
||||
seq,
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
// Fire-and-forget — caller doesn't need to await.
|
||||
let queue = queue.clone();
|
||||
tokio::spawn(async move {
|
||||
queue.publish_project_room_event(project_id, event).await;
|
||||
});
|
||||
}
|
||||
30
libs/room/src/service/patterns.rs
Normal file
30
libs/room/src/service/patterns.rs
Normal file
@ -0,0 +1,30 @@
|
||||
use std::sync::LazyLock;
|
||||
|
||||
/// Legacy: <user>uuid</user> or <user>username</user>
|
||||
static USER_MENTION_RE: LazyLock<regex_lite::Regex, fn() -> regex_lite::Regex> =
|
||||
LazyLock::new(|| regex_lite::Regex::new(r"<user>\s*([^<]+?)\s*</user>").unwrap());
|
||||
|
||||
/// Legacy: <mention type="..." id="...">label</mention>
|
||||
static MENTION_TAG_RE: LazyLock<regex_lite::Regex, fn() -> regex_lite::Regex> =
|
||||
LazyLock::new(|| {
|
||||
regex_lite::Regex::new(
|
||||
r#"<mention\s+type="([^"]+)"\s+id="([^"]+)"[^>]*>\s*([^<]*?)\s*</mention>"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
/// New format: @[type:id:label]
|
||||
static MENTION_BRACKET_RE: LazyLock<regex_lite::Regex, fn() -> regex_lite::Regex> =
|
||||
LazyLock::new(|| regex_lite::Regex::new(r"@\[([a-z]+):([^:\]]+):([^\]]+)\]").unwrap());
|
||||
|
||||
pub fn user_mention_re() -> &'static regex_lite::Regex {
|
||||
&USER_MENTION_RE
|
||||
}
|
||||
|
||||
pub fn mention_tag_re() -> &'static regex_lite::Regex {
|
||||
&MENTION_TAG_RE
|
||||
}
|
||||
|
||||
pub fn mention_bracket_re() -> &'static regex_lite::Regex {
|
||||
&MENTION_BRACKET_RE
|
||||
}
|
||||
49
libs/room/src/service/sequence.rs
Normal file
49
libs/room/src/service/sequence.rs
Normal file
@ -0,0 +1,49 @@
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room_message::{Column as RmCol, Entity as RoomMessage};
|
||||
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QuerySelect};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::RoomError;
|
||||
|
||||
pub async fn next_room_message_seq_internal(
|
||||
room_id: Uuid,
|
||||
db: &AppDatabase,
|
||||
cache: &AppCache,
|
||||
) -> Result<i64, RoomError> {
|
||||
let seq_key = format!("room:seq:{}", room_id);
|
||||
let mut conn = cache.conn().await.map_err(|e| {
|
||||
RoomError::Internal(format!("failed to get redis connection for seq: {}", e))
|
||||
})?;
|
||||
|
||||
let seq: i64 = redis::cmd("INCR")
|
||||
.arg(&seq_key)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.map_err(|e| RoomError::Internal(format!("seq INCR: {}", e)))?;
|
||||
|
||||
// DB reconciliation: only check every 1000 messages
|
||||
if seq % 1000 == 0 {
|
||||
let db_seq: Option<Option<Option<i64>>> = RoomMessage::find()
|
||||
.filter(RmCol::Room.eq(room_id))
|
||||
.select_only()
|
||||
.column_as(RmCol::Seq.max(), "max_seq")
|
||||
.into_tuple::<Option<Option<i64>>>()
|
||||
.one(db)
|
||||
.await?
|
||||
.map(|r| r);
|
||||
let db_seq = db_seq.flatten().flatten().unwrap_or(0);
|
||||
|
||||
if db_seq >= seq {
|
||||
let _: String = redis::cmd("SET")
|
||||
.arg(&seq_key)
|
||||
.arg(db_seq + 1)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.map_err(|e| RoomError::Internal(format!("seq SET: {}", e)))?;
|
||||
return Ok(db_seq + 1);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(seq)
|
||||
}
|
||||
329
libs/room/src/service/workers.rs
Normal file
329
libs/room/src/service/workers.rs
Normal file
@ -0,0 +1,329 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use db::cache::AppCache;
|
||||
use db::database::AppDatabase;
|
||||
use models::rooms::room;
|
||||
use queue::{AgentTaskEvent, MessageProducer};
|
||||
use sea_orm::EntityTrait;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::connection::{
|
||||
extract_get_redis, make_persist_fn, DedupCache, PersistFn, RoomConnectionManager,
|
||||
};
|
||||
|
||||
/// Callback type for sending push notifications.
|
||||
pub type PushNotificationFn =
|
||||
Arc<dyn Fn(Uuid, String, Option<String>, Option<String>) + Send + Sync>;
|
||||
|
||||
pub async fn start_workers(
|
||||
db: AppDatabase,
|
||||
_cache: AppCache,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
queue: MessageProducer,
|
||||
redis_url: String,
|
||||
dedup_cache: DedupCache,
|
||||
_task_service: Option<Arc<agent::TaskService>>,
|
||||
_max_concurrent_workers: Option<usize>,
|
||||
mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
|
||||
) -> anyhow::Result<()> {
|
||||
let rooms: Vec<room::Model> = room::Entity::find().all(&db).await?;
|
||||
let room_ids: Vec<uuid::Uuid> = rooms.iter().map(|r| r.id).collect();
|
||||
let project_ids: Vec<uuid::Uuid> = rooms
|
||||
.iter()
|
||||
.map(|r| r.project)
|
||||
.collect::<std::collections::HashSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let task_project_ids = project_ids.clone();
|
||||
|
||||
tracing::info!(
|
||||
room_count = room_ids.len(),
|
||||
project_count = project_ids.len(),
|
||||
"starting room workers"
|
||||
);
|
||||
|
||||
let persist_fn: PersistFn = make_persist_fn(db.clone(), room_manager.metrics.clone(), dedup_cache.clone());
|
||||
|
||||
let get_redis: Arc<dyn Fn() -> queue::worker::RedisFuture + Send + Sync> =
|
||||
extract_get_redis(queue.clone());
|
||||
|
||||
let worker_room_ids = room_ids.clone();
|
||||
let worker_shutdown = shutdown_rx.resubscribe();
|
||||
let worker_handle = tokio::spawn({
|
||||
let get_redis = get_redis.clone();
|
||||
let persist_fn = persist_fn.clone();
|
||||
async move {
|
||||
queue::start_worker(worker_room_ids, get_redis, persist_fn, worker_shutdown).await;
|
||||
}
|
||||
});
|
||||
|
||||
let manager = room_manager.clone();
|
||||
let redis_url_clone = redis_url.clone();
|
||||
|
||||
let mut handles: Vec<_> = room_ids
|
||||
.into_iter()
|
||||
.map(|room_id| {
|
||||
let manager = manager.clone();
|
||||
let redis_url = redis_url_clone.clone();
|
||||
let shutdown_rx = shutdown_rx.resubscribe();
|
||||
tokio::spawn(async move {
|
||||
crate::connection::subscribe_room_events(
|
||||
redis_url,
|
||||
manager,
|
||||
room_id,
|
||||
shutdown_rx,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let project_handles: Vec<_> = project_ids
|
||||
.into_iter()
|
||||
.map(|project_id| {
|
||||
let manager = manager.clone();
|
||||
let redis_url = redis_url_clone.clone();
|
||||
let shutdown_rx = shutdown_rx.resubscribe();
|
||||
tokio::spawn(async move {
|
||||
crate::connection::subscribe_project_room_events(
|
||||
redis_url,
|
||||
manager,
|
||||
project_id,
|
||||
shutdown_rx,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
handles.extend(project_handles);
|
||||
|
||||
let task_handles: Vec<_> = task_project_ids
|
||||
.into_iter()
|
||||
.map(|project_id| {
|
||||
let manager = manager.clone();
|
||||
let redis_url = redis_url_clone.clone();
|
||||
let shutdown_rx = shutdown_rx.resubscribe();
|
||||
tokio::spawn(async move {
|
||||
crate::connection::subscribe_task_events_fn(
|
||||
redis_url,
|
||||
manager,
|
||||
project_id,
|
||||
shutdown_rx,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
handles.extend(task_handles);
|
||||
|
||||
let cleanup_handle = {
|
||||
let manager = room_manager.clone();
|
||||
let db = db.clone();
|
||||
let dedup_cache = dedup_cache.clone();
|
||||
let mut cleanup_shutdown = shutdown_rx.resubscribe();
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(300));
|
||||
interval.tick().await;
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = interval.tick() => {
|
||||
manager.cleanup_rate_limit().await;
|
||||
crate::connection::cleanup_dedup_cache(&dedup_cache);
|
||||
if let Ok(rooms) = room::Entity::find().all(&db).await {
|
||||
let room_ids: Vec<_> = rooms.iter().map(|r| r.id).collect();
|
||||
let project_ids: Vec<_> = rooms.iter().map(|r| r.project).collect();
|
||||
manager.metrics.cleanup_stale_rooms(&room_ids).await;
|
||||
manager.prune_stale_rooms(&room_ids).await;
|
||||
manager.prune_stale_projects(&project_ids).await;
|
||||
}
|
||||
}
|
||||
_ = cleanup_shutdown.recv() => {
|
||||
tracing::info!("cleanup task shutting down");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
};
|
||||
handles.push(cleanup_handle);
|
||||
|
||||
let _ = shutdown_rx.recv().await;
|
||||
|
||||
tracing::info!("room workers shutting down");
|
||||
|
||||
for h in handles {
|
||||
let _ = h.abort();
|
||||
}
|
||||
let _ = worker_handle.await;
|
||||
|
||||
tracing::info!("room workers stopped");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn spawn_agent_task<F, Fut>(
|
||||
project_id: Uuid,
|
||||
agent_type: models::agent_task::AgentType,
|
||||
input: String,
|
||||
task_service: Arc<agent::TaskService>,
|
||||
queue: MessageProducer,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
worker_semaphore: Arc<tokio::sync::Semaphore>,
|
||||
execute: F,
|
||||
) -> anyhow::Result<i64>
|
||||
where
|
||||
F: FnOnce(i64, Arc<agent::TaskService>) -> Fut + Send + 'static,
|
||||
Fut: std::future::Future<Output = Result<String, String>> + Send,
|
||||
{
|
||||
let task = task_service
|
||||
.create(project_id, input, agent_type)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("create task failed: {}", e))?;
|
||||
|
||||
let task_id = task.id;
|
||||
|
||||
let started_event = AgentTaskEvent {
|
||||
task_id,
|
||||
project_id,
|
||||
parent_id: task.parent_id,
|
||||
event: "started".to_string(),
|
||||
message: None,
|
||||
output: None,
|
||||
error: None,
|
||||
status: models::agent_task::TaskStatus::Running.to_string(),
|
||||
timestamp: Utc::now(),
|
||||
};
|
||||
queue
|
||||
.publish_agent_task_event(project_id, started_event)
|
||||
.await;
|
||||
|
||||
let _ = task_service.start(task_id).await;
|
||||
|
||||
let queue_clone = queue.clone();
|
||||
let room_manager_clone = room_manager.clone();
|
||||
let semaphore = worker_semaphore.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _permit = semaphore.acquire().await.expect("semaphore closed");
|
||||
|
||||
let result = execute(task_id, task_service.clone()).await;
|
||||
|
||||
let event = match result {
|
||||
Ok(output) => {
|
||||
let _ = task_service.complete(task_id, &output).await;
|
||||
AgentTaskEvent {
|
||||
task_id,
|
||||
project_id,
|
||||
parent_id: None,
|
||||
event: "done".to_string(),
|
||||
message: None,
|
||||
output: Some(output),
|
||||
error: None,
|
||||
status: models::agent_task::TaskStatus::Done.to_string(),
|
||||
timestamp: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = task_service.fail(task_id, &err).await;
|
||||
AgentTaskEvent {
|
||||
task_id,
|
||||
project_id,
|
||||
parent_id: None,
|
||||
event: "failed".to_string(),
|
||||
message: None,
|
||||
output: None,
|
||||
error: Some(err),
|
||||
status: models::agent_task::TaskStatus::Failed.to_string(),
|
||||
timestamp: chrono::Utc::now(),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
queue_clone
|
||||
.publish_agent_task_event(project_id, event.clone())
|
||||
.await;
|
||||
room_manager_clone.broadcast_agent_task(project_id, event).await;
|
||||
tracing::info!(task_id = task_id, project_id = %project_id, "agent task finished");
|
||||
});
|
||||
|
||||
Ok(task_id)
|
||||
}
|
||||
|
||||
pub fn spawn_room_workers(
|
||||
room_id: uuid::Uuid,
|
||||
db: AppDatabase,
|
||||
room_manager: Arc<RoomConnectionManager>,
|
||||
queue: MessageProducer,
|
||||
redis_url: String,
|
||||
worker_semaphore: Arc<tokio::sync::Semaphore>,
|
||||
) {
|
||||
let persist_fn: PersistFn = make_persist_fn(
|
||||
db.clone(),
|
||||
room_manager.metrics.clone(),
|
||||
Arc::new(
|
||||
dashmap::DashMap::with_capacity_and_hasher(
|
||||
10000,
|
||||
Default::default(),
|
||||
),
|
||||
),
|
||||
);
|
||||
let get_redis: Arc<dyn Fn() -> queue::worker::RedisFuture + Send + Sync> =
|
||||
extract_get_redis(queue.clone());
|
||||
let manager = room_manager.clone();
|
||||
let redis_url_clone = redis_url.clone();
|
||||
let semaphore = worker_semaphore.clone();
|
||||
|
||||
let manager2 = room_manager.clone();
|
||||
let redis_url3 = redis_url.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let _permit = match semaphore.acquire_owned().await {
|
||||
Ok(p) => p,
|
||||
Err(_) => return,
|
||||
};
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::broadcast::channel::<()>(1);
|
||||
queue::room_worker_task(
|
||||
room_id,
|
||||
uuid::Uuid::new_v4().to_string(),
|
||||
get_redis,
|
||||
persist_fn,
|
||||
shutdown_rx,
|
||||
)
|
||||
.await;
|
||||
let _ = shutdown_tx.send(());
|
||||
});
|
||||
|
||||
tokio::spawn(async move {
|
||||
let shutdown_rx = manager.register_room(room_id).await;
|
||||
crate::connection::subscribe_room_events(
|
||||
redis_url_clone,
|
||||
manager.clone(),
|
||||
room_id,
|
||||
shutdown_rx,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
tokio::spawn(async move {
|
||||
let project_id = {
|
||||
let room = room::Entity::find_by_id(room_id)
|
||||
.one(&db)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
match room {
|
||||
Some(r) => r.project,
|
||||
None => return,
|
||||
}
|
||||
};
|
||||
let shutdown_rx = manager2.register_project(project_id).await;
|
||||
crate::connection::subscribe_project_room_events(
|
||||
redis_url3,
|
||||
manager2,
|
||||
project_id,
|
||||
shutdown_rx,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
@ -223,6 +223,9 @@ pub struct RoomMessageResponse {
|
||||
pub in_reply_to: Option<Uuid>,
|
||||
pub content: String,
|
||||
pub content_type: String,
|
||||
/// Accumulated AI reasoning/thinking text.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking_content: Option<String>,
|
||||
pub edited_at: Option<DateTime<Utc>>,
|
||||
pub send_at: DateTime<Utc>,
|
||||
pub revoked: Option<DateTime<Utc>>,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user