ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

feat: 添加媒体下载功能，优化消息发送逻辑，记录发送信息

Browse Source
This commit is contained in:
ooodc 2026-05-06 14:42:46 +08:00
parent 597881f72e
commit 32690cb792
2 changed files with 141 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

139
src/channels/wechat.rs
View File

@ -1,5 +1,6 @@
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
@ -89,6 +90,76 @@ impl WechatChannel {
}),
}
}
fn default_media_dir() -> PathBuf {
let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
home.join(".picobot").join("media").join("wechat")
}
fn build_download_filename(
media_type: &str,
file_name: Option<&str>,
format: Option<&str>,
) -> String {
if let Some(file_name) = file_name {
let sanitized: String = file_name
.chars()
.map(|ch| match ch {
'/' | '\\' | ':' | '*' | '?' | '"' | '<' | '>' | '|' => '_',
_ => ch,
})
.collect();
if !sanitized.trim().is_empty() {
return format!("{}_{}", uuid::Uuid::new_v4(), sanitized);
}
}
let ext = match (media_type, format) {
("image", _) => "jpg",
("video", _) => "mp4",
("voice", Some(fmt)) if !fmt.trim().is_empty() => fmt,
("voice", _) => "bin",
_ => "bin",
};
format!("{}_{}.{}", media_type, uuid::Uuid::new_v4(), ext)
}
async fn download_inbound_media(
bot: Arc<WeChatBot>,
msg: wechatbot::IncomingMessage,
) -> Result<Vec<MediaItem>, ChannelError> {
let Some(downloaded) = bot.download(&msg).await.map_err(|error| {
ChannelError::Other(format!("WeChat media download failed: {}", error))
})? else {
return Ok(Vec::new());
};
let media_dir = Self::default_media_dir();
tokio::fs::create_dir_all(&media_dir)
.await
.map_err(|error| ChannelError::Other(format!("Failed to create WeChat media dir: {}", error)))?;
let filename = Self::build_download_filename(
&downloaded.media_type,
downloaded.file_name.as_deref(),
downloaded.format.as_deref(),
);
let file_path = media_dir.join(&filename);
tokio::fs::write(&file_path, downloaded.data)
.await
.map_err(|error| ChannelError::Other(format!("Failed to write WeChat media file: {}", error)))?;
tracing::info!(filename = %filename, media_type = %downloaded.media_type, "Downloaded WeChat media");
let mut media_item = MediaItem::new(
file_path.to_string_lossy().to_string(),
downloaded.media_type,
);
media_item.mime_type = mime_guess::from_path(&file_path)
.first_raw()
.map(ToOwned::to_owned);
Ok(vec![media_item])
}
}
#[async_trait]
@ -109,6 +180,7 @@ impl Channel for WechatChannel {
let channel_name = self.name.clone();
let allow_from = self.config.allow_from.clone();
let bus_for_handler = bus.clone();
let bot_for_handler = self.bot.clone();
self.bot
.on_message(Box::new(move |msg| {
let sender_id = msg.user_id.clone();
@ -120,27 +192,38 @@ impl Channel for WechatChannel {
return;
}
let msg = msg.clone();
let timestamp = msg
.timestamp
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
let mut metadata = HashMap::new();
metadata.insert("context_token".to_string(), msg.context_token().to_string());
let inbound = InboundMessage {
channel: channel_name.clone(),
sender_id: sender_id.clone(),
chat_id: sender_id,
content: msg.text.clone(),
timestamp,
media: Vec::new(),
metadata,
forwarded_metadata: HashMap::new(),
};
let bus = bus_for_handler.clone();
let bot = bot_for_handler.clone();
let channel_name_for_publish = channel_name.clone();
tokio::spawn(async move {
let media = match Self::download_inbound_media(bot, msg.clone()).await {
Ok(media) => media,
Err(error) => {
tracing::error!(error = %error, "Failed to download WeChat inbound media");
Vec::new()
}
};
let mut metadata = HashMap::new();
metadata.insert("context_token".to_string(), msg.context_token().to_string());
let inbound = InboundMessage {
channel: channel_name_for_publish,
sender_id: sender_id.clone(),
chat_id: sender_id,
content: msg.text.clone(),
timestamp,
media,
metadata,
forwarded_metadata: HashMap::new(),
};
if let Err(error) = bus.publish_inbound(inbound).await {
tracing::error!(error = %error, "Failed to publish WeChat inbound message");
}
@ -202,6 +285,12 @@ impl Channel for WechatChannel {
self.bot.send(&msg.chat_id, &text).await.map_err(|error| {
ChannelError::SendError(format!("WeChat text send failed: {}", error))
})?;
tracing::info!(
channel = %self.name,
chat_id = %msg.chat_id,
content_len = text.len(),
"WeChat text message sent"
);
text_sent = true;
}
@ -215,6 +304,13 @@ impl Channel for WechatChannel {
self.bot.send_media(&msg.chat_id, content).await.map_err(|error| {
ChannelError::SendError(format!("WeChat media send failed: {}", error))
})?;
tracing::info!(
channel = %self.name,
chat_id = %msg.chat_id,
media_type = %media.media_type,
media_path = %media.path,
"WeChat media message sent"
);
}
if text.is_empty() && msg.media.is_empty() {
@ -234,6 +330,21 @@ mod tests {
use super::*;
use tempfile::NamedTempFile;
#[test]
fn build_download_filename_preserves_file_name() {
let filename = WechatChannel::build_download_filename("file", Some("README.md"), None);
assert!(filename.ends_with("_README.md"));
}
#[test]
fn build_download_filename_adds_voice_extension_when_missing_name() {
let filename = WechatChannel::build_download_filename("voice", None, Some("silk"));
assert!(filename.starts_with("voice_"));
assert!(filename.ends_with(".silk"));
}
#[test]
fn media_to_send_content_maps_image() {
let file = NamedTempFile::new().unwrap();

16
src/gateway/session_message_sender.rs
View File

@ -42,6 +42,7 @@ impl SessionMessageSender for BusSessionMessageSender {
.is_some();
if let Some(text) = request.text.filter(|value| !value.trim().is_empty()) {
let content_len = text.len();
self.bus
.publish_outbound(OutboundMessage::assistant(
channel_name.to_string(),
@ -52,10 +53,18 @@ impl SessionMessageSender for BusSessionMessageSender {
))
.await?;
published_messages += 1;
tracing::info!(
channel = %channel_name,
chat_id = %chat_id,
content_len = content_len,
"Published session text message to outbound bus"
);
}
let attachment_count = request.attachments.len();
for attachment in request.attachments {
let media_path = attachment.path.clone();
let media_type = attachment.media_type.clone();
let mut outbound = OutboundMessage::assistant(
channel_name.to_string(),
chat_id.to_string(),
@ -66,6 +75,13 @@ impl SessionMessageSender for BusSessionMessageSender {
outbound.media = vec![attachment];
self.bus.publish_outbound(outbound).await?;
published_messages += 1;
tracing::info!(
channel = %channel_name,
chat_id = %chat_id,
media_type = %media_type,
media_path = %media_path,
"Published session attachment to outbound bus"
);
}
Ok(SessionSendOutcome {