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:
ZhenYi 2026-04-26 13:10:42 +08:00
parent b4b5538447
commit f5e3da35b0
24 changed files with 2077 additions and 1761 deletions

View File

@ -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,

View File

@ -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,

View File

@ -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),

View File

@ -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())
}
}

View File

@ -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,

View File

@ -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,

View File

@ -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,

View File

@ -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

View 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()))
}

View 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)
}

View 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;
}
}
});
}

View 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;
}
}
});
}

View 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;
});
}

View 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;
});
}

View 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())
}
}
}

View 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()
}

View 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
}

View 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(())
}
}

View 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;
});
}

View 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
}

View 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)
}

View 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;
});
}

View File

@ -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>>,