refactor: route gateway messages concurrently
This commit is contained in:
parent
0c835db380
commit
515a07ec1f
@ -43,13 +43,13 @@ This file is the operational contract for coding agents working in this reposito
|
||||
### Core Data Flow
|
||||
|
||||
```
|
||||
Channel → MessageBus.inbound → Gateway processor → SessionManager → per-session worker → AgentLoop
|
||||
Channel → MessageBus.inbound → Gateway inbound router/lane → SessionManager → per-session worker → AgentLoop
|
||||
↑ │
|
||||
└── Channel ← OutboundDispatcher ← per-conversation lane ← MessageBus.outbound
|
||||
|
||||
AgentLoop → TurnEvent → Session TurnController → latest TurnSnapshot → DeliveryCoordinator → per-turn TurnSink → Channel
|
||||
|
||||
WebSocket/Channel → MessageBus.control → Gateway processor → SessionManager (dialog operations)
|
||||
WebSocket/Channel → MessageBus.control → Gateway control router → SessionManager (dialog operations)
|
||||
Scheduler → SessionManager scheduled execution → AgentLoop → Scheduler delivery policy → SessionManager/MessageBus
|
||||
```
|
||||
|
||||
|
||||
@ -39,7 +39,7 @@ Gateway 启动时先从配置目录 `.env`、workspace `.env` 和既有进程环
|
||||
flowchart LR
|
||||
External[CLI / Feishu] --> Channels[channels]
|
||||
Channels -->|InboundMessage| Bus[MessageBus]
|
||||
Bus --> Processor[Gateway message processor]
|
||||
Bus --> Processor[Gateway inbound/control routers]
|
||||
Processor --> Sessions[SessionManager]
|
||||
Sessions --> Agent[AgentLoop]
|
||||
Agent --> Providers[LLM providers]
|
||||
@ -79,9 +79,9 @@ flowchart LR
|
||||
|
||||
`MessageBus` 包含三条容量相同的 Tokio MPSC 队列:
|
||||
|
||||
- `inbound`:Channel → Gateway message processor。
|
||||
- `inbound`:Channel → Gateway inbound router。
|
||||
- `outbound`:Session/Tool → `OutboundDispatcher`。
|
||||
- `control`:WebSocket/Channel → Gateway message processor,用于 dialog 操作。
|
||||
- `control`:WebSocket/Channel → Gateway control router,用于 dialog 操作。
|
||||
|
||||
### 普通消息
|
||||
|
||||
@ -89,7 +89,7 @@ flowchart LR
|
||||
sequenceDiagram
|
||||
participant C as Channel
|
||||
participant B as MessageBus
|
||||
participant G as Message processor
|
||||
participant G as Inbound router
|
||||
participant S as SessionManager
|
||||
participant W as Per-session worker
|
||||
participant A as AgentLoop / Provider
|
||||
@ -120,7 +120,8 @@ sequenceDiagram
|
||||
|
||||
关键语义:
|
||||
|
||||
- Gateway 的主消息处理循环不等待模型完成;普通消息进入对应 session worker 后立即返回 `AgentProcessing`。
|
||||
- Gateway inbound router 按 `(channel, chat_id)` 使用容量 32 的短生命周期 lane 保持入口顺序,不同聊天可并发路由;普通消息进入对应 session worker 后立即返回 `AgentProcessing`。
|
||||
- `/stop` 绕过同聊天的 inbound lane,直接使正在运行的 worker/Turn 失效;其他 slash command 仍在聊天 lane 内有序执行。
|
||||
- 每个 session 有一条容量为 32 的队列,同一 session 串行处理,不同 session 的 worker 可并发执行。
|
||||
- 队列满时明确拒绝新消息,不允许无界积压。
|
||||
- Slash command 不进入 Agent 队列,由 `SessionManager` 直接执行,因此 `/stop` 等控制操作不会排在长模型调用之后。
|
||||
@ -154,7 +155,7 @@ sequenceDiagram
|
||||
|
||||
### Control 消息
|
||||
|
||||
WebSocket dialog 操作通过 `ControlMessage` 携带一次性回复通道。Gateway 在统一 message processor 中调用 `SessionManager`,再将 `SessionEvent` 回传给发起者。Bus 只承载消息,不解释操作。
|
||||
WebSocket dialog 操作通过 `ControlMessage` 携带一次性回复通道。Gateway control router 以最多 64 个并发受监督任务调用 `SessionManager`,再将 `SessionEvent` 回传给发起者;慢 control 不阻塞其他聊天的 inbound 路由。Bus 只承载消息,不解释操作。
|
||||
|
||||
TUI 的历史回放同样走 control 队列:`get_session_history` 先校验 session 属于当前客户端 scope,再由 SessionManager 从 Storage 读取最近消息。单次查询限制为 1–2000 条,TUI 默认请求最近 1000 条;迟到的历史响应只有在目标仍是当前 dialog 时才允许更新界面。
|
||||
|
||||
@ -212,7 +213,7 @@ SessionManager 负责组装会话上下文:系统提示、Skills、召回的 K
|
||||
|
||||
## 7. 后台任务与生命周期
|
||||
|
||||
`TaskSupervisor` 是 Gateway 内部后台任务的统一所有者。message processor、outbound dispatcher、scheduler、session workers、Turn delivery、outbound lanes、通知消费者和子 Agent 后台任务都应通过它注册。
|
||||
`TaskSupervisor` 是 Gateway 内部后台任务的统一所有者。inbound/control routers、inbound lanes、outbound dispatcher、scheduler、session workers、Turn delivery、outbound lanes、通知消费者和子 Agent 后台任务都应通过它注册。
|
||||
|
||||
两种注册方式:
|
||||
|
||||
@ -260,7 +261,7 @@ WebUI/TUI 文件字节通过受鉴权的 HTTP 接口流式传输,WebSocket 只
|
||||
3. 初始化 SQLite、MemoryManager、MessageBus 和 SessionManager;Scheduler 启用时幂等创建默认日常维护巡检。
|
||||
4. 注册内置工具、渠道、MCP 工具和 Cron 工具。
|
||||
5. 启动所有 Channel。
|
||||
6. 通过 TaskSupervisor 启动 message processor、dispatcher 和 scheduler。
|
||||
6. 通过 TaskSupervisor 启动 inbound/control routers、dispatcher 和 scheduler。
|
||||
7. 注册 WebUI 静态资源、管理 API 与聊天 WebSocket 路由。
|
||||
8. 绑定 Axum listener,开始接收请求。
|
||||
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
pub mod auth;
|
||||
pub mod http;
|
||||
mod router;
|
||||
pub mod uploads;
|
||||
pub mod ws;
|
||||
|
||||
@ -8,8 +9,7 @@ use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use crate::bus::{ControlMessage, MessageBus, OutboundDispatcher};
|
||||
use crate::channels::base::ChannelError;
|
||||
use crate::bus::{MessageBus, OutboundDispatcher};
|
||||
use crate::channels::{ChannelManager, CliChatChannel};
|
||||
use crate::config::{Config, ensure_workspace_dir, expand_path};
|
||||
use crate::delivery::{ConversationWriteLocks, DeliveryCoordinator, TurnDeliveryService};
|
||||
@ -276,80 +276,7 @@ impl GatewayState {
|
||||
}
|
||||
});
|
||||
|
||||
// Spawn unified message processor
|
||||
// This handles both inbound AI messages and control messages in one loop
|
||||
self.task_supervisor.spawn("message-processor", async move {
|
||||
tracing::info!("Message processor started");
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// Inbound: AI message flow
|
||||
inbound = bus.consume_inbound() => {
|
||||
let Some(inbound) = inbound else {
|
||||
tracing::warn!("Message processor stopping because inbound bus closed");
|
||||
break;
|
||||
};
|
||||
match session_manager.handle_message(
|
||||
&inbound.channel,
|
||||
&inbound.sender_id,
|
||||
&inbound.chat_id,
|
||||
&inbound.content,
|
||||
inbound.media,
|
||||
inbound.forwarded_metadata.clone(),
|
||||
).await {
|
||||
Ok(crate::session::session::HandleResult::AgentResponse(content)) => {
|
||||
let outbound = crate::bus::OutboundMessage {
|
||||
channel: inbound.channel.clone(),
|
||||
chat_id: inbound.chat_id.clone(),
|
||||
content,
|
||||
reply_to: None,
|
||||
media: vec![],
|
||||
metadata: inbound.forwarded_metadata,
|
||||
delivery: None,
|
||||
};
|
||||
if let Err(e) = bus.publish_outbound(outbound).await {
|
||||
tracing::error!(error = %e, "Failed to publish outbound");
|
||||
}
|
||||
}
|
||||
Ok(crate::session::session::HandleResult::CommandOutput(content)) => {
|
||||
let mut metadata = inbound.forwarded_metadata;
|
||||
metadata.insert("_type".to_string(), "command".to_string());
|
||||
let outbound = crate::bus::OutboundMessage {
|
||||
channel: inbound.channel.clone(),
|
||||
chat_id: inbound.chat_id.clone(),
|
||||
content,
|
||||
reply_to: None,
|
||||
media: vec![],
|
||||
metadata,
|
||||
delivery: None,
|
||||
};
|
||||
if let Err(e) = bus.publish_outbound(outbound).await {
|
||||
tracing::error!(error = %e, "Failed to publish outbound");
|
||||
}
|
||||
}
|
||||
Ok(crate::session::session::HandleResult::AgentProcessing) => {
|
||||
// Agent is processing in background; response will be
|
||||
// sent via bus directly from the spawned task.
|
||||
// The select loop remains free to handle subsequent
|
||||
// messages (including slash commands).
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "Failed to handle message");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Control: session management operations
|
||||
msg = bus.consume_control() => {
|
||||
let Some(msg) = msg else {
|
||||
tracing::warn!("Message processor stopping because control bus closed");
|
||||
break;
|
||||
};
|
||||
Self::handle_control_message(&session_manager, msg).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
router::spawn_message_routers(bus.clone(), session_manager, self.task_supervisor.clone());
|
||||
|
||||
// Spawn outbound dispatcher
|
||||
let dispatcher = OutboundDispatcher::new(
|
||||
@ -379,112 +306,6 @@ impl GatewayState {
|
||||
tracing::info!("Scheduler background task spawned");
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle control messages (session management operations)
|
||||
async fn handle_control_message(session_manager: &SessionManager, msg: ControlMessage) {
|
||||
use crate::session::{SessionCommand::*, SessionEvent};
|
||||
|
||||
let reply_tx = msg.reply_tx;
|
||||
let result: Result<SessionEvent, ChannelError> = match msg.op {
|
||||
CreateDialog {
|
||||
channel,
|
||||
chat_id,
|
||||
title,
|
||||
} => session_manager
|
||||
.create_dialog(&channel, &chat_id, title.as_deref())
|
||||
.await
|
||||
.map(|(session_id, title)| SessionEvent::DialogCreated { session_id, title })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
ListDialogs {
|
||||
channel,
|
||||
chat_id,
|
||||
include_archived,
|
||||
} => session_manager
|
||||
.list_dialogs(&channel, &chat_id, include_archived)
|
||||
.await
|
||||
.map(|(dialogs, current_dialog_id)| SessionEvent::DialogList {
|
||||
dialogs,
|
||||
current_dialog_id,
|
||||
})
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
GetCurrentDialog { channel, chat_id } => session_manager
|
||||
.get_current_dialog(&channel, &chat_id)
|
||||
.await
|
||||
.map(|session_id| SessionEvent::CurrentDialog { session_id })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
SwitchDialog {
|
||||
channel,
|
||||
chat_id,
|
||||
dialog_id,
|
||||
} => session_manager
|
||||
.switch_dialog(&channel, &chat_id, &dialog_id)
|
||||
.await
|
||||
.map(|session_id| SessionEvent::DialogSwitched { session_id })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
GetDialogHistory { session_id, limit } => session_manager
|
||||
.get_dialog_history(&session_id, limit)
|
||||
.await
|
||||
.map(|messages| SessionEvent::DialogHistory {
|
||||
session_id,
|
||||
messages,
|
||||
})
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
GetTaskPlan { session_id } => session_manager
|
||||
.get_task_plan(&session_id)
|
||||
.await
|
||||
.map(|plan| SessionEvent::TaskPlan { session_id, plan })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
RenameDialog { session_id, title } => session_manager
|
||||
.rename_dialog(&session_id, &title)
|
||||
.await
|
||||
.map(|()| SessionEvent::DialogRenamed { session_id, title })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
ArchiveDialog { session_id } => session_manager
|
||||
.archive_dialog(&session_id)
|
||||
.await
|
||||
.map(|()| SessionEvent::DialogArchived { session_id })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
DeleteDialog { session_id } => session_manager
|
||||
.delete_dialog(&session_id)
|
||||
.await
|
||||
.map(|()| SessionEvent::DialogDeleted { session_id })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
ClearHistory { session_id } => session_manager
|
||||
.clear_dialog_history(&session_id)
|
||||
.await
|
||||
.map(|()| SessionEvent::HistoryCleared { session_id })
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
GetSlashCommands {
|
||||
channel: _,
|
||||
chat_id: _,
|
||||
} => {
|
||||
let commands = session_manager.get_slash_commands().to_vec();
|
||||
Ok(SessionEvent::SlashCommandsList { commands })
|
||||
}
|
||||
ExecuteSlashCommand {
|
||||
command,
|
||||
args,
|
||||
channel,
|
||||
chat_id,
|
||||
current_session_id,
|
||||
} => session_manager
|
||||
.execute_slash_command(
|
||||
&command,
|
||||
args.as_deref(),
|
||||
&channel,
|
||||
&chat_id,
|
||||
current_session_id.as_ref(),
|
||||
)
|
||||
.await
|
||||
.map(|(new_id, msg)| SessionEvent::SlashCommandExecuted {
|
||||
new_session_id: new_id,
|
||||
message: msg,
|
||||
})
|
||||
.map_err(|e| ChannelError::Other(e.to_string())),
|
||||
};
|
||||
|
||||
let _ = reply_tx.send(result).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run(
|
||||
|
||||
424
src/gateway/router.rs
Normal file
424
src/gateway/router.rs
Normal file
@ -0,0 +1,424 @@
|
||||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::{Semaphore, mpsc};
|
||||
|
||||
use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage};
|
||||
use crate::channels::ChannelError;
|
||||
use crate::channels::parse_slash_command;
|
||||
use crate::session::{SessionCommand, SessionEvent, SessionManager};
|
||||
use crate::task_supervisor::TaskSupervisor;
|
||||
|
||||
const INBOUND_LANE_CAPACITY: usize = 32;
|
||||
const INBOUND_LANE_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
const CONTROL_MAX_IN_FLIGHT: usize = 64;
|
||||
|
||||
pub(super) fn spawn_message_routers(
|
||||
bus: Arc<MessageBus>,
|
||||
session_manager: Arc<SessionManager>,
|
||||
supervisor: TaskSupervisor,
|
||||
) {
|
||||
spawn_inbound_router(bus.clone(), session_manager.clone(), supervisor.clone());
|
||||
spawn_control_router(bus, session_manager, supervisor);
|
||||
}
|
||||
|
||||
fn spawn_inbound_router(
|
||||
bus: Arc<MessageBus>,
|
||||
session_manager: Arc<SessionManager>,
|
||||
supervisor: TaskSupervisor,
|
||||
) {
|
||||
let lane_supervisor = supervisor.clone();
|
||||
supervisor.spawn("inbound-router", async move {
|
||||
tracing::info!(lane_capacity = INBOUND_LANE_CAPACITY, "Inbound router started");
|
||||
let mut lanes: HashMap<String, mpsc::Sender<InboundMessage>> = HashMap::new();
|
||||
let mut messages_seen = 0_u64;
|
||||
|
||||
while let Some(inbound) = bus.consume_inbound().await {
|
||||
messages_seen = messages_seen.wrapping_add(1);
|
||||
if messages_seen.is_multiple_of(128) {
|
||||
lanes.retain(|_, sender| !sender.is_closed());
|
||||
}
|
||||
|
||||
// Stop must be able to invalidate a running worker even when an
|
||||
// earlier slow slash command occupies this conversation's lane.
|
||||
if is_priority_stop(&inbound.content) {
|
||||
let request_bus = bus.clone();
|
||||
let request_manager = session_manager.clone();
|
||||
let task_name = format!("inbound-stop:{}:{}", inbound.channel, inbound.chat_id);
|
||||
if !lane_supervisor.spawn(task_name, async move {
|
||||
process_inbound(request_bus, request_manager, inbound).await;
|
||||
}) {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let key = conversation_key(&inbound.channel, &inbound.chat_id);
|
||||
let mut sender = lanes.get(&key).cloned();
|
||||
if sender.as_ref().is_none_or(mpsc::Sender::is_closed) {
|
||||
let (new_sender, receiver) = mpsc::channel(INBOUND_LANE_CAPACITY);
|
||||
if !spawn_inbound_lane(
|
||||
&lane_supervisor,
|
||||
bus.clone(),
|
||||
session_manager.clone(),
|
||||
inbound.channel.clone(),
|
||||
inbound.chat_id.clone(),
|
||||
receiver,
|
||||
) {
|
||||
tracing::warn!("Inbound router is stopping");
|
||||
break;
|
||||
}
|
||||
lanes.insert(key.clone(), new_sender.clone());
|
||||
sender = Some(new_sender);
|
||||
}
|
||||
|
||||
let Some(sender) = sender else {
|
||||
tracing::error!("Inbound lane creation did not produce a sender");
|
||||
continue;
|
||||
};
|
||||
match sender.try_send(inbound) {
|
||||
Ok(()) => {}
|
||||
Err(mpsc::error::TrySendError::Full(inbound)) => {
|
||||
tracing::warn!(channel = %inbound.channel, chat_id = %inbound.chat_id, "Inbound conversation lane is full");
|
||||
publish_command_output(
|
||||
&bus,
|
||||
inbound,
|
||||
"当前对话入口队列已满,请稍后重试。".to_string(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Closed(inbound)) => {
|
||||
// The lane may have exited on its idle boundary between the
|
||||
// closed check and try_send. Recreate it once without
|
||||
// dropping this input.
|
||||
let (new_sender, receiver) = mpsc::channel(INBOUND_LANE_CAPACITY);
|
||||
if !spawn_inbound_lane(
|
||||
&lane_supervisor,
|
||||
bus.clone(),
|
||||
session_manager.clone(),
|
||||
inbound.channel.clone(),
|
||||
inbound.chat_id.clone(),
|
||||
receiver,
|
||||
) {
|
||||
break;
|
||||
}
|
||||
lanes.insert(key, new_sender.clone());
|
||||
if let Err(error) = new_sender.try_send(inbound) {
|
||||
tracing::error!(error = %error, "Failed to enqueue input into replacement lane");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
tracing::warn!("Inbound router stopped because inbound bus closed");
|
||||
});
|
||||
}
|
||||
|
||||
fn spawn_inbound_lane(
|
||||
supervisor: &TaskSupervisor,
|
||||
bus: Arc<MessageBus>,
|
||||
session_manager: Arc<SessionManager>,
|
||||
channel: String,
|
||||
chat_id: String,
|
||||
receiver: mpsc::Receiver<InboundMessage>,
|
||||
) -> bool {
|
||||
supervisor.spawn(format!("inbound-lane:{channel}:{chat_id}"), async move {
|
||||
run_ordered_lane(receiver, INBOUND_LANE_IDLE_TIMEOUT, move |inbound| {
|
||||
process_inbound(bus.clone(), session_manager.clone(), inbound)
|
||||
})
|
||||
.await;
|
||||
})
|
||||
}
|
||||
|
||||
async fn run_ordered_lane<T, F, Fut>(
|
||||
mut receiver: mpsc::Receiver<T>,
|
||||
idle_timeout: Duration,
|
||||
mut handler: F,
|
||||
) where
|
||||
T: Send + 'static,
|
||||
F: FnMut(T) -> Fut,
|
||||
Fut: Future<Output = ()>,
|
||||
{
|
||||
loop {
|
||||
let item = match tokio::time::timeout(idle_timeout, receiver.recv()).await {
|
||||
Ok(Some(item)) => item,
|
||||
Ok(None) | Err(_) => break,
|
||||
};
|
||||
handler(item).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn process_inbound(
|
||||
bus: Arc<MessageBus>,
|
||||
session_manager: Arc<SessionManager>,
|
||||
inbound: InboundMessage,
|
||||
) {
|
||||
let result = session_manager
|
||||
.handle_message(
|
||||
&inbound.channel,
|
||||
&inbound.sender_id,
|
||||
&inbound.chat_id,
|
||||
&inbound.content,
|
||||
inbound.media.clone(),
|
||||
inbound.forwarded_metadata.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(crate::session::session::HandleResult::AgentResponse(content)) => {
|
||||
publish_assistant_output(&bus, inbound, content).await;
|
||||
}
|
||||
Ok(crate::session::session::HandleResult::CommandOutput(content)) => {
|
||||
publish_command_output(&bus, inbound, content).await;
|
||||
}
|
||||
Ok(crate::session::session::HandleResult::AgentProcessing) => {}
|
||||
Err(error) => {
|
||||
tracing::error!(channel = %inbound.channel, chat_id = %inbound.chat_id, error = %error, "Failed to handle inbound message");
|
||||
publish_command_output(&bus, inbound, "消息处理失败,请稍后重试。".to_string()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn publish_assistant_output(bus: &MessageBus, inbound: InboundMessage, content: String) {
|
||||
publish_output(bus, inbound, content, false).await;
|
||||
}
|
||||
|
||||
async fn publish_command_output(bus: &MessageBus, inbound: InboundMessage, content: String) {
|
||||
publish_output(bus, inbound, content, true).await;
|
||||
}
|
||||
|
||||
async fn publish_output(bus: &MessageBus, inbound: InboundMessage, content: String, command: bool) {
|
||||
let mut metadata = inbound.forwarded_metadata;
|
||||
if command {
|
||||
metadata.insert("_type".to_string(), "command".to_string());
|
||||
}
|
||||
let outbound = OutboundMessage {
|
||||
channel: inbound.channel,
|
||||
chat_id: inbound.chat_id,
|
||||
content,
|
||||
reply_to: None,
|
||||
media: vec![],
|
||||
metadata,
|
||||
delivery: None,
|
||||
};
|
||||
if let Err(error) = bus.publish_outbound(outbound).await {
|
||||
tracing::error!(error = %error, "Failed to publish routed outbound message");
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_control_router(
|
||||
bus: Arc<MessageBus>,
|
||||
session_manager: Arc<SessionManager>,
|
||||
supervisor: TaskSupervisor,
|
||||
) {
|
||||
let request_supervisor = supervisor.clone();
|
||||
supervisor.spawn("control-router", async move {
|
||||
tracing::info!(
|
||||
max_in_flight = CONTROL_MAX_IN_FLIGHT,
|
||||
"Control router started"
|
||||
);
|
||||
let permits = Arc::new(Semaphore::new(CONTROL_MAX_IN_FLIGHT));
|
||||
loop {
|
||||
let permit = match permits.clone().acquire_owned().await {
|
||||
Ok(permit) => permit,
|
||||
Err(_) => break,
|
||||
};
|
||||
let Some(message) = bus.consume_control().await else {
|
||||
break;
|
||||
};
|
||||
let manager = session_manager.clone();
|
||||
if !request_supervisor.spawn("control-request", async move {
|
||||
let _permit = permit;
|
||||
handle_control_message(&manager, message).await;
|
||||
}) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
tracing::warn!("Control router stopped because control bus closed");
|
||||
});
|
||||
}
|
||||
|
||||
async fn handle_control_message(session_manager: &SessionManager, message: ControlMessage) {
|
||||
use SessionCommand::*;
|
||||
|
||||
let reply_tx = message.reply_tx;
|
||||
let result: Result<SessionEvent, ChannelError> = match message.op {
|
||||
CreateDialog {
|
||||
channel,
|
||||
chat_id,
|
||||
title,
|
||||
} => session_manager
|
||||
.create_dialog(&channel, &chat_id, title.as_deref())
|
||||
.await
|
||||
.map(|(session_id, title)| SessionEvent::DialogCreated { session_id, title })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
ListDialogs {
|
||||
channel,
|
||||
chat_id,
|
||||
include_archived,
|
||||
} => session_manager
|
||||
.list_dialogs(&channel, &chat_id, include_archived)
|
||||
.await
|
||||
.map(|(dialogs, current_dialog_id)| SessionEvent::DialogList {
|
||||
dialogs,
|
||||
current_dialog_id,
|
||||
})
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
GetCurrentDialog { channel, chat_id } => session_manager
|
||||
.get_current_dialog(&channel, &chat_id)
|
||||
.await
|
||||
.map(|session_id| SessionEvent::CurrentDialog { session_id })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
SwitchDialog {
|
||||
channel,
|
||||
chat_id,
|
||||
dialog_id,
|
||||
} => session_manager
|
||||
.switch_dialog(&channel, &chat_id, &dialog_id)
|
||||
.await
|
||||
.map(|session_id| SessionEvent::DialogSwitched { session_id })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
GetDialogHistory { session_id, limit } => session_manager
|
||||
.get_dialog_history(&session_id, limit)
|
||||
.await
|
||||
.map(|messages| SessionEvent::DialogHistory {
|
||||
session_id,
|
||||
messages,
|
||||
})
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
GetTaskPlan { session_id } => session_manager
|
||||
.get_task_plan(&session_id)
|
||||
.await
|
||||
.map(|plan| SessionEvent::TaskPlan { session_id, plan })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
RenameDialog { session_id, title } => session_manager
|
||||
.rename_dialog(&session_id, &title)
|
||||
.await
|
||||
.map(|()| SessionEvent::DialogRenamed { session_id, title })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
ArchiveDialog { session_id } => session_manager
|
||||
.archive_dialog(&session_id)
|
||||
.await
|
||||
.map(|()| SessionEvent::DialogArchived { session_id })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
DeleteDialog { session_id } => session_manager
|
||||
.delete_dialog(&session_id)
|
||||
.await
|
||||
.map(|()| SessionEvent::DialogDeleted { session_id })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
ClearHistory { session_id } => session_manager
|
||||
.clear_dialog_history(&session_id)
|
||||
.await
|
||||
.map(|()| SessionEvent::HistoryCleared { session_id })
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
GetSlashCommands { .. } => Ok(SessionEvent::SlashCommandsList {
|
||||
commands: session_manager.get_slash_commands().to_vec(),
|
||||
}),
|
||||
ExecuteSlashCommand {
|
||||
command,
|
||||
args,
|
||||
channel,
|
||||
chat_id,
|
||||
current_session_id,
|
||||
} => session_manager
|
||||
.execute_slash_command(
|
||||
&command,
|
||||
args.as_deref(),
|
||||
&channel,
|
||||
&chat_id,
|
||||
current_session_id.as_ref(),
|
||||
)
|
||||
.await
|
||||
.map(
|
||||
|(new_session_id, message)| SessionEvent::SlashCommandExecuted {
|
||||
new_session_id,
|
||||
message,
|
||||
},
|
||||
)
|
||||
.map_err(|error| ChannelError::Other(error.to_string())),
|
||||
};
|
||||
|
||||
let _ = reply_tx.send(result).await;
|
||||
}
|
||||
|
||||
fn conversation_key(channel: &str, chat_id: &str) -> String {
|
||||
format!("{channel}\0{chat_id}")
|
||||
}
|
||||
|
||||
fn is_priority_stop(content: &str) -> bool {
|
||||
parse_slash_command(content).is_some_and(|(command, _)| command == "stop")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::HashSet;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
#[test]
|
||||
fn only_stop_bypasses_a_conversation_lane() {
|
||||
assert!(is_priority_stop("/stop"));
|
||||
assert!(is_priority_stop(" /stop "));
|
||||
assert!(!is_priority_stop("/compact"));
|
||||
assert!(!is_priority_stop("normal message"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_keys_do_not_alias() {
|
||||
let keys = HashSet::from([
|
||||
conversation_key("a", "bc"),
|
||||
conversation_key("ab", "c"),
|
||||
conversation_key("a", "bd"),
|
||||
]);
|
||||
assert_eq!(keys.len(), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn slow_conversation_lane_does_not_block_another_lane() {
|
||||
let (slow_tx, slow_rx) = mpsc::channel(2);
|
||||
let (fast_tx, fast_rx) = mpsc::channel(2);
|
||||
let slow_started = Arc::new(Notify::new());
|
||||
let release_slow = Arc::new(Notify::new());
|
||||
let fast_finished = Arc::new(Notify::new());
|
||||
|
||||
let slow_task = tokio::spawn({
|
||||
let slow_started = slow_started.clone();
|
||||
let release_slow = release_slow.clone();
|
||||
async move {
|
||||
run_ordered_lane(slow_rx, Duration::from_secs(1), move |_| {
|
||||
let slow_started = slow_started.clone();
|
||||
let release_slow = release_slow.clone();
|
||||
async move {
|
||||
slow_started.notify_one();
|
||||
release_slow.notified().await;
|
||||
}
|
||||
})
|
||||
.await;
|
||||
}
|
||||
});
|
||||
let fast_task = tokio::spawn({
|
||||
let fast_finished = fast_finished.clone();
|
||||
async move {
|
||||
run_ordered_lane(fast_rx, Duration::from_secs(1), move |_| {
|
||||
let fast_finished = fast_finished.clone();
|
||||
async move { fast_finished.notify_one() }
|
||||
})
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
slow_tx.send("slow").await.unwrap();
|
||||
slow_started.notified().await;
|
||||
fast_tx.send("fast").await.unwrap();
|
||||
tokio::time::timeout(Duration::from_millis(100), fast_finished.notified())
|
||||
.await
|
||||
.expect("fast lane was blocked by unrelated slow lane");
|
||||
|
||||
release_slow.notify_one();
|
||||
drop(slow_tx);
|
||||
drop(fast_tx);
|
||||
slow_task.await.unwrap();
|
||||
fast_task.await.unwrap();
|
||||
}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user