xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化斜杠命令处理机制，增加/stop命令

Browse Source
This commit is contained in:
xiaoski 2026-05-24 18:12:18 +08:00
parent b84c6f85db
commit 22be6e404b
3 changed files with 393 additions and 148 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/gateway/mod.rs
View File

@ -223,6 +223,12 @@ impl GatewayState {
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");
}

3
src/scheduler/mod.rs
View File

@ -179,6 +179,9 @@ impl Scheduler {
let _ = self.storage.record_scheduled_job_run(&run).await;
}
Ok(HandleResult::AgentProcessing) => {
tracing::warn!(job_id = %job.id, "scheduler: unexpected AgentProcessing from cron — response sent via bus");
}
Err(e) => {
let error_str = e.to_string();
let run = JobRun {

406
src/session/session.rs
View File

@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::{Mutex, mpsc, oneshot};
use crate::bus::{ChatMessage, MediaItem, MessageSource, OutboundMessage, SourceKind};
use crate::mcp::get_mcp_status;
@ -18,6 +18,8 @@ pub enum HandleResult {
AgentResponse(String),
/// Command output to be sent as CommandExecuted
CommandOutput(String),
/// Agent processing spawned in background; response will be sent via bus
AgentProcessing,
}
use crate::channels::slash_command::parse_slash_command;
use crate::config::LLMProviderConfig;
@ -72,6 +74,21 @@ pub struct Session {
pub last_compressed_message_at: Option<i64>,
#[allow(dead_code)]
memory_manager: Arc<crate::memory::MemoryManager>,
/// Task queue for per-session serial agent processing
agent_tx: Option<mpsc::UnboundedSender<AgentTask>>,
/// Cancel signal for the currently executing agent task
current_cancel: Option<oneshot::Sender<()>>,
/// Monotonic counter to detect stale workers
worker_generation: u64,
}
/// A task to be processed by the per-session agent worker
struct AgentTask {
channel: String,
chat_id: String,
content: String,
media: Vec<MediaItem>,
}
impl Session {
@ -119,6 +136,9 @@ impl Session {
last_consolidated_at: None,
last_compressed_message_at: None,
memory_manager,
agent_tx: None,
current_cancel: None,
worker_generation: 0,
})
}
@ -251,6 +271,9 @@ impl Session {
last_consolidated_at: session_meta.last_consolidated_at,
last_compressed_message_at: session_meta.last_compressed_message_at,
memory_manager,
agent_tx: None,
current_cancel: None,
worker_generation: 0,
})
}
@ -806,6 +829,11 @@ pub static SLASH_COMMANDS: &[SlashCommand] = &[
description: "显示 MCP 服务状态和工具列表",
aliases: &["/mcp"],
},
SlashCommand {
name: "stop",
description: "停止当前正在执行的任务并清空消息队列",
aliases: &["/stop"],
},
];
impl SessionManager {
@ -1027,6 +1055,26 @@ impl SessionManager {
}).collect();
Ok((None, format!("MCP 服务：\n\n{}", lines.join("\n\n"))))
}
"stop" => {
let sid = current_session_id
.ok_or_else(|| AgentError::Other("no active session".to_string()))?;
let session = self.get_or_create_session(sid).await?;
let mut guard = session.lock().await;
let mut msgs: Vec<String> = Vec::new();
if guard.current_cancel.take().is_some() {
msgs.push("当前任务已发送停止信号。".to_string());
}
if guard.agent_tx.take().is_some() {
msgs.push("消息队列已清空。".to_string());
}
guard.worker_generation = guard.worker_generation.wrapping_add(1);
let resp = if msgs.is_empty() {
"没有正在执行的任务或队列。".to_string()
} else {
msgs.join(" ")
};
Ok((None, resp))
}
_ => Err(AgentError::Other(format!("未知命令：/{}。输入 /? 获取帮助。", cmd.name))),
}
}
@ -1348,14 +1396,86 @@ impl SessionManager {
};
}
// Normal message handling through LLM
let (notify_tx, mut notify_rx) = tokio::sync::mpsc::unbounded_channel();
// Spawn notification publisher — sends immediately when tools are detected
// Normal message: enqueue to per-session worker for serial processing.
let task = AgentTask {
channel: channel.to_string(),
chat_id: chat_id.to_string(),
content: content.to_string(),
media,
};
let session_clone = session.clone();
let unified_str = unified_id.to_string();
{
let bus = self.bus.clone();
let ch = channel.to_string();
let cid = chat_id.to_string();
let mut guard = session_clone.lock().await;
let needs_spawn =
guard.agent_tx.is_none() || guard.agent_tx.as_ref().is_some_and(|tx| tx.is_closed());
if needs_spawn {
guard.agent_tx = None;
guard.current_cancel = None;
guard.worker_generation = guard.worker_generation.wrapping_add(1);
let generation = guard.worker_generation;
let (tx, rx) = mpsc::unbounded_channel();
guard.agent_tx = Some(tx);
spawn_agent_worker(
rx,
session_clone.clone(),
self.bus.clone(),
self.memory_manager.clone(),
self.skills_loader.clone(),
generation,
unified_str.clone(),
);
}
if let Err(e) = guard.agent_tx.as_ref().unwrap().send(task) {
// Worker died after we just spawned it — respawn with the recovered task
let task = e.0;
guard.agent_tx = None;
guard.current_cancel = None;
guard.worker_generation = guard.worker_generation.wrapping_add(1);
let generation = guard.worker_generation;
let (tx, rx) = mpsc::unbounded_channel();
guard.agent_tx = Some(tx);
spawn_agent_worker(
rx,
session_clone.clone(),
self.bus.clone(),
self.memory_manager.clone(),
self.skills_loader.clone(),
generation,
unified_str.clone(),
);
guard.agent_tx.as_ref().unwrap().send(task).unwrap_or_else(|_| {
tracing::error!("Agent worker spawn+send failed irrecoverably");
});
}
}
Ok(HandleResult::AgentProcessing)
}).await
}
}
fn spawn_agent_worker(
mut task_rx: mpsc::UnboundedReceiver<AgentTask>,
session: Arc<Mutex<Session>>,
bus: Arc<MessageBus>,
memory_manager: Arc<crate::memory::MemoryManager>,
skills_loader: Arc<SkillsLoader>,
worker_gen: u64,
unified_str: String,
) {
tokio::spawn(async move {
let _scope = CURRENT_SOURCE_SESSION.scope(Some(unified_str), async {
while let Some(task) = task_rx.recv().await {
let task_chan = task.channel.clone();
let task_cid = task.chat_id.clone();
let (notify_tx, mut notify_rx) = mpsc::unbounded_channel();
// Spawn notification publisher
{
let bus = bus.clone();
let ch = task_chan.clone();
let cid = task_cid.clone();
tokio::spawn(async move {
while let Some(notif) = notify_rx.recv().await {
let mut metadata = HashMap::new();
@ -1374,30 +1494,51 @@ impl SessionManager {
}
// Phase 1: prepare data under session lock
let (agent, history, system_prompt) = {
let mut session_guard = session.lock().await;
let (agent, history_out, system_prompt_out, cancel_rx) = {
let mut guard = session.lock().await;
let media_refs: Vec<String> = media.iter().map(|m| m.path.clone()).collect();
#[cfg(debug_assertions)]
if !media_refs.is_empty() {
tracing::debug!(media_count = %media.len(), media_refs = ?media_refs, "Adding user message with media");
if guard.worker_generation != worker_gen {
return; // stale worker
}
let user_message = session_guard.create_user_message(content, media_refs);
session_guard.add_message(user_message, true).await
.map_err(|e| AgentError::Other(format!("persist error: {}", e)))?;
let media_refs: Vec<String> =
task.media.iter().map(|m| m.path.clone()).collect();
let user_message =
guard.create_user_message(&task.content, media_refs);
if let Err(e) = guard.add_message(user_message, true).await {
tracing::error!(error = %e, "Failed to persist user message");
let err_outbound = OutboundMessage {
channel: task_chan.clone(),
chat_id: task_cid.clone(),
content: "Failed to save your message, please try again."
.to_string(),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus.publish_outbound(err_outbound).await;
return;
}
let history = session_guard.get_history().to_vec();
let history_raw = guard.get_history().to_vec();
let skills_prompt = skills_loader.build_skills_prompt();
let skills_prompt = self.skills_loader.build_skills_prompt();
let memory_context = match self.memory_manager.recall(content, 5, Some(crate::memory::MemoryCategory::Knowledge), None).await {
Ok(entries) if !entries.is_empty() => {
Some(entries.iter()
let memory_context = match memory_manager
.recall(
&task.content,
5,
Some(crate::memory::MemoryCategory::Knowledge),
None,
)
.await
{
Ok(entries) if !entries.is_empty() => Some(
entries
.iter()
.map(|e| format!("- {}: {}", e.key, e.content))
.collect::<Vec<_>>()
.join("\n"))
}
.join("\n"),
),
Err(e) => {
tracing::warn!(error = %e, "Failed to fetch memory context");
None
@ -1405,101 +1546,196 @@ impl SessionManager {
_ => None,
};
let system_prompt = session_guard.build_system_prompt(&skills_prompt, memory_context.as_deref());
let system_prompt = guard
.build_system_prompt(&skills_prompt, memory_context.as_deref());
let result = session_guard.compressor
.compress_if_needed(history)
let result = guard
.compressor
.compress_if_needed(history_raw)
.await
.inspect_err(|e| {
tracing::warn!(error = %e, "Context compression failed in handle_message");
})?;
if result.created_timelines {
session_guard.last_compressed_message_at = Some(chrono::Utc::now().timestamp_millis());
.map(|r| {
if r.created_timelines {
guard.last_compressed_message_at =
Some(chrono::Utc::now().timestamp_millis());
}
let mut history = result.history;
r.history
})
.unwrap_or_else(|e| {
tracing::warn!(
error = %e,
"Context compression failed in worker"
);
guard.get_history().to_vec()
});
let mut history = result;
history.insert(0, ChatMessage::system(system_prompt.clone()));
let now = chrono::Utc::now().timestamp_millis();
session_guard.last_consolidated_at = Some(now);
if let Err(e) = session_guard.persist_session_meta().await {
tracing::warn!(error = %e, "Failed to persist consolidation timestamp");
guard.last_consolidated_at = Some(now);
let _ = guard.persist_session_meta().await;
let agent = match guard.create_agent_with_notify(notify_tx) {
Ok(a) => a,
Err(e) => {
tracing::error!(error = %e, "Failed to create agent");
let err_outbound = OutboundMessage {
channel: task_chan.clone(),
chat_id: task_cid.clone(),
content: "Agent creation failed, please try again."
.to_string(),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus.publish_outbound(err_outbound).await;
return;
}
};
let agent = session_guard.create_agent_with_notify(notify_tx)?;
(agent, history, system_prompt)
}; // session lock released — send_message can now lock freely
let (cancel_tx, cancel_rx) = oneshot::channel();
// Phase 2: LLM call (no session lock held)
let result = match agent.process(history).await {
if guard.worker_generation != worker_gen {
return; // /stop replaced us
}
guard.current_cancel = Some(cancel_tx);
(agent, history, system_prompt, cancel_rx)
}; // lock released
// Phase 2 + 3: LLM call with cancellation
let session2 = session.clone();
let bus2 = bus.clone();
let chan2 = task_chan.clone();
let cid2 = task_cid.clone();
let process_future = async move {
let result = match agent.process(history_out.clone()).await {
Ok(r) => r,
Err(AgentError::LlmError(ref msg))
if is_context_overflow_error(msg) =>
{
let retry_history = {
let mut session_guard = session.lock().await;
let new_window = crate::agent::ContextCompressor::parse_context_limit_from_error(msg)
.unwrap_or(session_guard.compressor_threshold());
let mut guard = session2.lock().await;
let new_window =
crate::agent::ContextCompressor::parse_context_limit_from_error(msg)
.unwrap_or(guard.compressor_threshold());
tracing::warn!(
new_window,
error = %msg,
"Context overflow in handle_message — retrying with tighter window"
"Context overflow in worker — retrying"
);
session_guard.compressor.set_context_window(new_window);
let raw = session_guard.get_history().to_vec();
let retry_result = session_guard.compressor.compress_if_needed(raw).await?;
if retry_result.created_timelines {
session_guard.last_compressed_message_at = Some(chrono::Utc::now().timestamp_millis());
if let Err(e) = session_guard.persist_session_meta().await {
tracing::warn!(error = %e, "Failed to persist compression marker on retry");
guard.compressor.set_context_window(new_window);
let raw = guard.get_history().to_vec();
let retry_result =
match guard.compressor.compress_if_needed(raw).await {
Ok(r) => r,
Err(e) => {
tracing::error!(error = %e, "Retry compression failed");
let err_outbound = OutboundMessage {
channel: chan2,
chat_id: cid2,
content: "Context overflow handling failed."
.to_string(),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
}
};
if retry_result.created_timelines {
guard.last_compressed_message_at =
Some(chrono::Utc::now().timestamp_millis());
let _ = guard.persist_session_meta().await;
}
let mut retry = retry_result.history;
retry.insert(0, ChatMessage::system(system_prompt.clone()));
retry.insert(
0,
ChatMessage::system(system_prompt_out.clone()),
);
retry
}; // lock released again for retry
agent.process(retry_history).await
.inspect_err(|e| {
tracing::error!(error = %e, "Agent retry after context compression failed");
})?
}
Err(e) => {
tracing::error!(error = %e, "Agent processing error — propagating to caller");
return Err(e);
},
};
// Phase 3: persist results under session lock
let response: String = {
let mut session_guard = session.lock().await;
for msg in result.emitted_messages {
session_guard.add_message(msg, true).await
.map_err(|e| AgentError::Other(format!("persist error: {}", e)))?;
match agent.process(retry_history).await {
Ok(r) => r,
Err(e) => {
tracing::error!(
error = %e,
"Agent retry after overflow failed"
);
let err_outbound = OutboundMessage {
channel: chan2,
chat_id: cid2,
content: format!("Processing error: {}", e),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
}
}
}
Err(e) => {
tracing::error!(error = %e, "Agent processing error");
let err_outbound = OutboundMessage {
channel: chan2,
chat_id: cid2,
content: format!("Processing error: {}", e),
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus2.publish_outbound(err_outbound).await;
return;
}
};
if session_guard.should_generate_title()
&& let Err(e) = session_guard.generate_title().await {
let response = {
let mut guard = session2.lock().await;
for msg in result.emitted_messages {
guard.add_message(msg, true).await.inspect_err(|e| {
tracing::error!(error = %e, "Failed to persist message")
}).ok();
}
if guard.should_generate_title()
&& let Err(e) = guard.generate_title().await
{
tracing::warn!("failed to generate title: {}", e);
}
result.final_response.content
};
#[cfg(debug_assertions)]
tracing::debug!(
channel = %channel,
chat_id = %chat_id,
response_len = %response.len(),
"Agent response received"
);
let outbound = OutboundMessage {
channel: chan2,
chat_id: cid2,
content: response,
reply_to: None,
media: vec![],
metadata: HashMap::new(),
};
let _ = bus2.publish_outbound(outbound).await;
};
Ok(HandleResult::AgentResponse(response))
}).await
tokio::select! {
() = process_future => {}
_ = cancel_rx => {
// cancelled — current_cancel already taken by /stop
}
}
/// Handle a message triggered by a scheduled cron job.
// Clean up
let mut guard = session.lock().await;
if guard.worker_generation == worker_gen {
guard.current_cancel = None;
}
}
}).await;
});
}
impl SessionManager {
///
/// Runs in a stateless manner: no session creation, no history persistence.
/// The cron system prompt instructs the LLM to deliver results via the