新增: df-im-gateway 网关骨架

This commit is contained in:
lxy
2026-10-05 00:35:00 +08:00
parent 91773da602
commit e492f03191
37 changed files with 4211 additions and 12 deletions
+26
View File
@@ -0,0 +1,26 @@
[package]
name = "df-im-gateway"
version = "0.1.0"
edition = "2021"
description = "IM 网关:飞书/企微/钉钉/Web 四端远程接入 DevFlow AI Chat(WS Commander 连 u-relay)"
[dependencies]
anyhow = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
rusqlite = { workspace = true }
tokio = { workspace = true }
async-trait = { workspace = true }
tracing = { workspace = true }
df-tunnel = { path = "../df-tunnel" }
# WS(与 df-tunnel 同款本地版本声明,workspace 无这几项)
tokio-tungstenite = { version = "0.23", features = ["native-tls"] }
futures-util = "0.3"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
# 飞书 SDK(长连接 + 归一化消息;channel 自动连带 ws+messaging)
larksuite-oapi-sdk-rs = { version = "0.3.12", features = ["channel"] }
[dev-dependencies]
tempfile = "3"
+251
View File
@@ -0,0 +1,251 @@
//! 平台适配器(P0 预留项 ④:trait 分层,统一 InboundEvent/OutboundAction)
//!
//! 网关核心只认识这两个类型,四端差异(飞书卡片/企微 Markdown/钉钉卡片/Web JSON)收敛在各 adapter。
//! P0 只实现 Web,飞书/企微/钉钉留 trait 位(P1/P2/P4 各自落地,不动核心)。
use async_trait::async_trait;
use tokio::sync::mpsc;
/// 平台入站事件(IM 平台 → 网关,已归一化)
#[derive(Debug, Clone)]
pub enum InboundEvent {
/// 用户发消息
Message {
platform: String,
chat_id: String,
user_id: String,
text: String,
},
/// 会话切换(/new 或菜单)
NewSession { platform: String, chat_id: String },
/// 审批按钮回调
Approval {
platform: String,
chat_id: String,
approval_id: String,
approved: bool,
},
}
/// 平台出站动作(网关 → IM 平台,已归一化)
#[derive(Debug, Clone)]
pub enum OutboundAction {
/// 发/更新 AI 流式文本(全量内容,适配器自选卡片或文本)
StreamText { chat_id: String, text: String, done: bool },
/// 发送审批卡片(同意/拒绝按钮)
ApprovalCard { chat_id: String, approval_id: String, title: String, detail: String },
/// 发系统提示
Notice { chat_id: String, text: String },
}
/// 平台适配器
#[async_trait]
pub trait PlatformAdapter: Send + Sync {
/// 平台名(feishu/wecom/dingtalk/web,与 session_map.platform 对应)。
/// 核心循环经 platform_index 注册表间接使用。
#[allow(dead_code)]
fn name(&self) -> &str;
/// 启动适配器(长连接/回调服务器/WS),入站事件从 events 通道流出。阻塞运行。
async fn run(&self, events: mpsc::Sender<InboundEvent>) -> anyhow::Result<()>;
/// 出站动作分发(网关核心调用)
async fn dispatch(&self, action: OutboundAction) -> anyhow::Result<()>;
/// 清理平台的流式消息状态(网关在新建会话映射时调用,兜底上轮异常中止的残留)。
/// 默认空实现:无流式覆写语义的平台(Web 直接推帧)无需关心。
async fn reset_stream(&self, _chat_id: &str) {}
}
/// Web Chat 适配器(网关直接托管浏览器 WS,无三方依赖)
pub mod web {
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt};
use tokio::sync::{mpsc, Mutex};
use tokio_tungstenite::tungstenite::Message;
use super::{InboundEvent, OutboundAction, PlatformAdapter};
/// 浏览器连接注册表(chat_id → ws 发送通道)
type WsRegistry = Arc<Mutex<HashMap<String, mpsc::Sender<Message>>>>;
/// Web 适配器:浏览器 WS 直连网关 0.0.0.0:9191(accept 协议见 handle_web_ws)
#[derive(Clone)]
pub struct WebAdapter {
registry: WsRegistry,
}
impl WebAdapter {
pub fn new() -> Self {
Self { registry: Arc::new(Mutex::new(HashMap::new())) }
}
}
impl Default for WebAdapter {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl PlatformAdapter for WebAdapter {
fn name(&self) -> &str {
"web"
}
async fn run(&self, events: mpsc::Sender<InboundEvent>) -> anyhow::Result<()> {
let listener = tokio::net::TcpListener::bind("0.0.0.0:9191").await?;
tracing::info!("[web] WS 服务监听 0.0.0.0:9191");
loop {
let (stream, _peer) = listener.accept().await?;
let registry = self.registry.clone();
let events = events.clone();
tokio::spawn(async move {
if let Err(e) = handle_web_ws(stream, registry, events).await {
tracing::warn!("[web] 浏览器连接异常退出: {e}");
}
});
}
}
async fn dispatch(&self, action: OutboundAction) -> anyhow::Result<()> {
match action {
OutboundAction::StreamText { chat_id, text, done } => {
let frame = serde_json::json!({
"type": if done { "stream_done" } else { "stream_delta" },
"text": text,
});
self.send_ws(&chat_id, &frame.to_string()).await
}
OutboundAction::ApprovalCard { chat_id, approval_id, title, detail } => {
let frame = serde_json::json!({
"type": "approval",
"approval_id": approval_id,
"title": title,
"detail": detail,
});
self.send_ws(&chat_id, &frame.to_string()).await
}
OutboundAction::Notice { chat_id, text } => {
let frame = serde_json::json!({ "type": "notice", "text": text });
self.send_ws(&chat_id, &frame.to_string()).await
}
}
}
}
impl WebAdapter {
async fn send_ws(&self, chat_id: &str, text: &str) -> anyhow::Result<()> {
let registry = self.registry.lock().await;
if let Some(tx) = registry.get(chat_id) {
tx.send(Message::Text(text.to_string())).await?;
}
Ok(())
}
}
/// 单个浏览器 WS 会话。
///
/// 协议(浏览器 → 网关):
/// ```json
/// 首帧绑定: {"type":"hello","chat_id":"...","user_id":"..."}
/// 消息: {"type":"message","text":"..."}
/// 重置会话: {"type":"new_session"}
/// 审批: {"type":"approval","approval_id":"...","approved":true|false}
/// ```
/// 网关 → 浏览器:stream_delta / stream_done / approval / notice(见 dispatch)。
async fn handle_web_ws(
stream: tokio::net::TcpStream,
registry: WsRegistry,
events: mpsc::Sender<InboundEvent>,
) -> anyhow::Result<()> {
let ws = tokio_tungstenite::accept_async(stream).await?;
let (mut sink, mut upstream) = ws.split();
// 首帧 hello 绑定 chat_id
let first = upstream.next().await.ok_or_else(|| anyhow::anyhow!("浏览器连接即断"))??;
let hello: serde_json::Value = match &first {
Message::Text(t) => serde_json::from_str(t)?,
_ => anyhow::bail!("首帧非文本"),
};
let chat_id = hello
.get("chat_id")
.and_then(|v| v.as_str())
.unwrap_or("web-anon")
.to_string();
let user_id = hello.get("user_id").and_then(|v| v.as_str()).unwrap_or("anon").to_string();
// 注册发送通道:出站帧经 ws_tx → forward 任务写 socket
let (ws_tx, mut ws_rx) = mpsc::channel::<Message>(64);
registry.lock().await.insert(chat_id.clone(), ws_tx);
tracing::info!("[web] 浏览器接入 chat_id={chat_id} user={user_id}");
let forward = tokio::spawn(async move {
while let Some(frame) = ws_rx.recv().await {
if sink.send(frame).await.is_err() {
break;
}
}
});
let result: anyhow::Result<()> = async {
while let Some(msg) = upstream.next().await {
match msg? {
Message::Text(t) => {
let v: serde_json::Value = serde_json::from_str(&t).unwrap_or_default();
match v.get("type").and_then(|x| x.as_str()).unwrap_or("") {
"message" => {
let text = v.get("text").and_then(|x| x.as_str()).unwrap_or("").to_string();
events
.send(InboundEvent::Message {
platform: "web".into(),
chat_id: chat_id.clone(),
user_id: user_id.clone(),
text,
})
.await
.map_err(|_| anyhow::anyhow!("核心入站通道关闭"))?;
}
"new_session" => {
events
.send(InboundEvent::NewSession {
platform: "web".into(),
chat_id: chat_id.clone(),
})
.await
.map_err(|_| anyhow::anyhow!("核心入站通道关闭"))?;
}
"approval" => {
let approval_id =
v.get("approval_id").and_then(|x| x.as_str()).unwrap_or("").to_string();
let approved = v.get("approved").and_then(|x| x.as_bool()).unwrap_or(false);
events
.send(InboundEvent::Approval {
platform: "web".into(),
chat_id: chat_id.clone(),
approval_id,
approved,
})
.await
.map_err(|_| anyhow::anyhow!("核心入站通道关闭"))?;
}
_ => {}
}
}
Message::Close(_) => break,
_ => {}
}
}
Ok(())
}
.await;
registry.lock().await.remove(&chat_id);
forward.abort();
tracing::info!("[web] 浏览器断开 chat_id={chat_id}");
result
}
}
+188
View File
@@ -0,0 +1,188 @@
//! WS Commander 客户端(连 u-relay /ws/commander)
//!
//! 与桌面端 df-tunnel(Agent 角色)对称:网关作为 Commander 发 TunnelCommand、收 TunnelEvent。
//! 协议类型直接复用 df_tunnel::events(编译期一致)。
//!
//! 连接生命周期:hello {kind:"commander", channel_id, token} → (relay 校验) → 双向消息。
//! 断线重连:指数退避(1/2/4/8/16/30s 封顶,与桌面端 df-tunnel 同款策略)。
//!
//! 注:relay 对同 channel 多 Commander 的投递行为(P0 必验项 ①)待实测——
//! 过渡期约定同 channel 单活跃控制端。
use std::time::Duration;
use futures_util::{SinkExt, StreamExt};
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use tracing::{info, warn};
use df_tunnel::events::{Hello, HelloKind, RelayControlError};
use df_tunnel::{TunnelCommand, TunnelEvent, TunnelMessage};
/// 指数退避序列(秒,封顶 30)
fn backoff_secs(attempt: u32) -> u64 {
[1, 2, 4, 8, 16, 30]
.get(attempt.min(5) as usize)
.copied()
.unwrap_or(30)
}
/// 启动 Commander 连接循环(阻塞直到 shutdown)。重连在内部循环,指令/事件双向透传。
///
/// * `relay_url` — ws://host:9180/ws/commander
/// * `channel_id` — 配对的桌面端 channel
/// * `token` — U_RELAY_TOKEN 同源
/// * `ev_tx` — 事件出口(桌面端 → 网关核心)
/// * `cmd_rx` — 指令入口(网关核心 → 桌面端)
/// * `shutdown` — 关闭信号
pub async fn run_commander(
relay_url: String,
channel_id: String,
token: String,
ev_tx: mpsc::Sender<TunnelEvent>,
mut cmd_rx: mpsc::Receiver<TunnelCommand>,
shutdown: tokio::sync::watch::Receiver<bool>,
) -> anyhow::Result<()> {
let mut attempt: u32 = 0;
loop {
if *shutdown.borrow() {
info!("[commander] 收到关闭信号,退出");
return Ok(());
}
match connect_once(&relay_url, &channel_id, &token, &ev_tx, &mut cmd_rx, shutdown.clone()).await {
Ok(()) => {
info!("[commander] 连接正常关闭");
return Ok(());
}
Err(e) => {
let wait = backoff_secs(attempt);
attempt += 1;
warn!("[commander] 连接断开({e}),{}s 后重连(第 {attempt} 次)", wait);
let mut shutdown_w = shutdown.clone();
tokio::select! {
_ = tokio::time::sleep(Duration::from_secs(wait)) => {}
_ = shutdown_w.changed() => return Ok(()),
}
}
}
}
}
/// 单次连接会话(hello 握手 + 消息循环),断开/出错返回 Err 触发外层重连
async fn connect_once(
relay_url: &str,
channel_id: &str,
token: &str,
ev_tx: &mpsc::Sender<TunnelEvent>,
cmd_rx: &mut mpsc::Receiver<TunnelCommand>,
mut shutdown: tokio::sync::watch::Receiver<bool>,
) -> anyhow::Result<()> {
let (ws, _resp) = tokio_tungstenite::connect_async(relay_url).await?;
let (mut sink, mut stream) = ws.split();
// hello 握手(kind=commander;字段与 u-relay::relay::Hello 严格一致)
let hello = Hello {
kind: HelloKind::Commander,
channel_id: channel_id.to_string(),
token: token.to_string(),
};
sink.send(Message::Text(serde_json::to_string(&hello)?)).await?;
// 等待 hello_ack(首条文本帧;control error 帧则报错)
let ack = tokio::time::timeout(Duration::from_secs(10), stream.next()).await
.map_err(|_| anyhow::anyhow!("hello_ack 10s 超时"))?
.ok_or_else(|| anyhow::anyhow!("连接在握手期关闭"))?;
match ack? {
Message::Text(t) => {
// relay 控制帧:{"kind":"control","error":"..."} → 握手被拒
if let Ok(ctrl_err) = serde_json::from_str::<RelayControlError>(&t) {
if ctrl_err.kind == "control" {
anyhow::bail!("握手被拒: {}", ctrl_err.error);
}
}
info!("[commander] hello_ack: channel={channel_id} 握手成功");
}
other => anyhow::bail!("hello_ack 非文本帧: {other:?}"),
}
// 心跳:30s ping(relay 90s 读死线)
let mut ping = tokio::time::interval(Duration::from_secs(30));
ping.tick().await; // 首个立即返回,跳过
// 消息循环
loop {
tokio::select! {
_ = shutdown.changed() => {
let _ = sink.close().await;
return Ok(());
}
_ = ping.tick() => {
sink.send(Message::Ping(vec![])).await?;
}
cmd = cmd_rx.recv() => {
let Some(cmd) = cmd else {
// 核心侧通道关闭:正常收尾
let _ = sink.close().await;
return Ok(());
};
let frame = TunnelMessage::Command(cmd);
sink.send(Message::Text(serde_json::to_string(&frame)?)).await?;
}
msg = stream.next() => {
match msg {
Some(Ok(Message::Text(t))) => {
match serde_json::from_str::<TunnelMessage>(&t) {
Ok(TunnelMessage::Event(e)) => {
if ev_tx.send(e).await.is_err() {
return Err(anyhow::anyhow!("核心事件通道关闭"));
}
}
Ok(TunnelMessage::Command(_)) => {
warn!("[commander] commander 收到 Command 帧(协议错向,忽略)");
}
Err(_) => warn!("[commander] 无法解析帧(非 TunnelMessage,忽略): {} 字节", t.len()),
}
}
Some(Ok(Message::Ping(p))) => { let _ = sink.send(Message::Pong(p)).await; }
Some(Ok(_)) => {}
Some(Err(e)) => return Err(e.into()),
None => return Err(anyhow::anyhow!("stream 结束")),
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn backoff_sequence() {
assert_eq!(backoff_secs(0), 1);
assert_eq!(backoff_secs(1), 2);
assert_eq!(backoff_secs(2), 4);
assert_eq!(backoff_secs(3), 8);
assert_eq!(backoff_secs(4), 16);
assert_eq!(backoff_secs(5), 30);
assert_eq!(backoff_secs(99), 30, "封顶 30s");
}
#[test]
fn tunnel_message_roundtrip() {
let cmd = TunnelMessage::Command(TunnelCommand::Send {
conv_id: "c1".into(),
content: "你好".into(),
});
let json = serde_json::to_string(&cmd).unwrap();
assert!(json.contains("\"command\""));
let back: TunnelMessage = serde_json::from_str(&json).unwrap();
match back {
TunnelMessage::Command(TunnelCommand::Send { conv_id, content }) => {
assert_eq!(conv_id, "c1");
assert_eq!(content, "你好");
}
_ => panic!("roundtrip 失败"),
}
}
}
+143
View File
@@ -0,0 +1,143 @@
//! 网关配置(env 加载,凭证不进 git)
//!
//! 配置来源:环境变量 > 默认值(gateway.toml 文件加载为后续增强,当前未实现)。
//! u-relay token 与桌面端 df-tunnel 同源(env U_RELAY_TOKEN)。
use serde::Deserialize;
/// 网关全局配置
#[derive(Debug, Clone, Deserialize)]
pub struct GatewayConfig {
/// u-relay WS 地址(生产 u-relay.1216.top,本地 http://127.0.0.1:9180)
pub relay_url: String,
/// relay 配对 token(env U_RELAY_TOKEN 同源)
pub relay_token: String,
/// channel 路由表:IM 侧默认走 default_channel;按前缀路由的多桌面端配置(预留项 ⑥)
#[serde(default)]
pub channels: Vec<ChannelRoute>,
/// 飞书应用凭证(未配置则该平台禁用)
#[serde(default)]
pub feishu: Option<FeishuConfig>,
/// TextDelta 节流窗口毫秒(默认 800,IM 频控友好)
#[serde(default = "default_throttle_ms")]
pub throttle_ms: u64,
/// 用户白名单(空=放行所有,生产部署务必显式配置)
///
/// 注:历史注释「空=拒绝所有」与 main.rs 实现(空=放行)矛盾,走查 2026-10-04 定为
/// 空=放行——否则首次启动零配置时飞书/Web 全部拒绝,连 banner 警告都没有,反而
/// 更难发现配置缺失。生产风险由 main.rs 启动 banner 常驻警告 + 白名单缺失提示兜底。
#[serde(default)]
pub whitelist: Vec<String>,
}
/// channel 路由项:IM 指令前缀 → 桌面端 channel
#[derive(Debug, Clone, Deserialize)]
pub struct ChannelRoute {
/// 指令前缀(如 "/dev2 xxx" 路由到 chan-home);空串=默认 channel
pub prefix: String,
/// u-relay channel_id(桌面端 df-tunnel 配置的同一 channel)
pub channel_id: String,
}
/// 飞书企业自建应用凭证
#[derive(Debug, Clone, Deserialize)]
pub struct FeishuConfig {
pub app_id: String,
pub app_secret: String,
/// 长连接模式事件订阅(免公网回调;当前恒 true,保留显式开关供回调模式扩展)
#[allow(dead_code)]
pub long_connection: bool,
}
fn default_throttle_ms() -> u64 {
800
}
impl GatewayConfig {
/// 从 env + 可选 toml 文件构建。缺失关键项(relay_url/token)返回 Err。
pub fn load() -> anyhow::Result<Self> {
let relay_url = std::env::var("GATEWAY_RELAY_URL")
.unwrap_or_else(|_| "ws://127.0.0.1:9180/ws/commander".to_string());
let relay_token = std::env::var("GATEWAY_RELAY_TOKEN")
.or_else(|_| std::env::var("U_RELAY_TOKEN"))
.map_err(|_| anyhow::anyhow!("缺少 GATEWAY_RELAY_TOKEN / U_RELAY_TOKEN 环境变量"))?;
let feishu = match (
std::env::var("FEISHU_APP_ID"),
std::env::var("FEISHU_APP_SECRET"),
) {
(Ok(app_id), Ok(app_secret)) => Some(FeishuConfig {
app_id,
app_secret,
long_connection: true,
}),
_ => None,
};
let whitelist = std::env::var("GATEWAY_WHITELIST")
.map(|v| v.split(',').map(|s| s.trim().to_string()).filter(|s| !s.is_empty()).collect())
.unwrap_or_default();
let default_channel = std::env::var("GATEWAY_DEFAULT_CHANNEL").unwrap_or_default();
let channels = if default_channel.is_empty() {
Vec::new()
} else {
vec![ChannelRoute { prefix: String::new(), channel_id: default_channel }]
};
Ok(Self {
relay_url,
relay_token,
channels,
feishu,
throttle_ms: default_throttle_ms(),
whitelist,
})
}
/// 查指令应路由到的 channel(前缀匹配,空前缀=兜底默认)
pub fn route(&self, command_text: &str) -> Option<&str> {
// 精确前缀优先,空前缀兜底
self.channels
.iter()
.find(|c| !c.prefix.is_empty() && command_text.starts_with(&c.prefix))
.map(|c| c.channel_id.as_str())
.or_else(|| self.channels.iter().find(|c| c.prefix.is_empty()).map(|c| c.channel_id.as_str()))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn route_prefix_priority() {
let cfg = GatewayConfig {
relay_url: String::new(),
relay_token: String::new(),
channels: vec![
ChannelRoute { prefix: "/dev2".into(), channel_id: "chan-home".into() },
ChannelRoute { prefix: String::new(), channel_id: "chan-work".into() },
],
feishu: None,
throttle_ms: 800,
whitelist: vec![],
};
assert_eq!(cfg.route("/dev2 修个 bug"), Some("chan-home"));
assert_eq!(cfg.route("普通指令"), Some("chan-work"));
assert_eq!(cfg.route("/dev2"), Some("chan-home"));
}
#[test]
fn route_empty_when_no_channels() {
let cfg = GatewayConfig {
relay_url: String::new(),
relay_token: String::new(),
channels: vec![],
feishu: None,
throttle_ms: 800,
whitelist: vec![],
};
assert_eq!(cfg.route("任何"), None);
}
}
+350
View File
@@ -0,0 +1,350 @@
//! 飞书适配器(P2):长连接事件 + 归一化消息 + 卡片审批
//!
//! 依赖 larksuite-oapi-sdk-rs features=["channel"](自动连带 ws+messaging)。
//! 长连接:飞书←→网关出站直连,不经 u-relay(u-relay 仅穿 NAT 段到桌面端)。
//!
//! 验证记录(2026-10-04,SDK 0.3.12 源码核对):
//! - 卡片回调:ChannelBuilder::on_card_action(NormalizedCardAction)
//! 内部走 dispatcher.on_card_action_trigger(channel/handler.rs:112),
//! 即新版 card.action.trigger 事件通道,长连接可收(遗留验证项 ✅)。
//! - 发卡片:SendInput.card(String) → INTERACTIVE 类型(messaging.rs:752)。
//! - 流式更新:channel_messaging().send() 拿 message_id 后 edit_text() 覆写全文。
use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use larksuite_oapi_sdk_rs::card::v1::{
Action, ActionComponent, Button, ButtonType, Card, CardDocument, Div, Element, Header,
Markdown,
};
use larksuite_oapi_sdk_rs::channel::{Channel, ChannelPolicy};
use larksuite_oapi_sdk_rs::{EventDispatcher, LarkClient, LarkError, RequestOption};
use tokio::sync::mpsc::Sender;
use crate::adapters::{InboundEvent, OutboundAction, PlatformAdapter};
use crate::config::FeishuConfig;
/// 飞书适配器:channel 归一化消息 + 卡片审批,单飞书应用单实例
#[derive(Clone)]
pub struct FeishuAdapter {
client: LarkClient,
/// 流式会话注册表:chat_id → 已发出的流式消息 message_id(用于 edit_text 覆写)
streams: Arc<tokio::sync::Mutex<HashMap<String, String>>>,
}
impl FeishuAdapter {
pub fn new(cfg: &FeishuConfig) -> anyhow::Result<Self> {
let client = LarkClient::builder(&cfg.app_id, &cfg.app_secret).build()?;
Ok(Self {
client,
streams: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
})
}
/// 构造失败时的离线哨兵:client 用占位凭证(不启动 run,dispatch 记日志即弃)
pub fn offline() -> Self {
let client = LarkClient::builder("offline", "offline").build()
.expect("占位凭证构造 LarkClient 不发网络请求,必成功");
Self {
client,
streams: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
}
}
}
#[async_trait]
impl PlatformAdapter for FeishuAdapter {
fn name(&self) -> &str {
"feishu"
}
/// 兜底清理流式残留(走查 #3):桌面端异常中止时无 Completed 事件,
/// streams entry 永驻导致下轮首帧错覆写旧消息。核心循环在新建会话映射时调用。
async fn reset_stream(&self, chat_id: &str) {
self.streams.lock().await.remove(chat_id);
}
async fn run(&self, events: Sender<InboundEvent>) -> anyhow::Result<()> {
// 长连接模式:空 verification_token/encrypt_key(免公网回调)
let dispatcher = EventDispatcher::new("", "");
let client = self.client.clone();
// 入站事件通道复制两份,分别给消息/卡片回调闭包
let events_for_msg = events.clone();
let events_for_card = events;
let channel: Channel<'static> = {
// Channel<'a> 借用 client;用 Box::leak 提升为 'static(适配器生命周期即进程生命周期)
let leaked: &'static LarkClient = Box::leak(Box::new(client));
Channel::builder(leaked, dispatcher)
.policy(ChannelPolicy::default().require_mention(false))
.on_message(move |msg| {
let events = events_for_msg.clone();
async move {
handle_inbound_message(events, msg).await;
Ok(())
}
})
.on_card_action(move |action| {
let events = events_for_card.clone();
async move {
handle_card_action(events, action).await;
Ok(())
}
})
.build()
};
tracing::info!("[feishu] 长连接启动(免公网回调)");
channel.start().await.map_err(|e: LarkError| anyhow::anyhow!("飞书长连接退出: {e}"))?;
Ok(())
}
async fn dispatch(&self, action: OutboundAction) -> anyhow::Result<()> {
let opt = &RequestOption::default();
match action {
OutboundAction::StreamText { chat_id, text, done } => {
// 两段式:锁内只取/写 message_id(纯内存操作),网络 IO 在锁外执行。
// 修复走查 #2:同 chat 仍串行(顺序语义),不同 chat 不再被同一把锁挡住。
let existing_id: Option<String> = {
let streams = self.streams.lock().await;
streams.get(&chat_id).cloned()
};
if let Some(message_id) = existing_id {
let res = self.client.channel_messaging().edit_text(&message_id, &text, opt).await;
if let Err(e) = res {
tracing::warn!("[feishu] 流式更新失败(消息可能已被撤回/超时): {e},回退新发");
self.streams.lock().await.remove(&chat_id);
}
} else {
// 首帧:发新文本消息,记录 message_id 供后续覆写
let input = larksuite_oapi_sdk_rs::channel::SendInput {
receive_id: Some(chat_id.clone()),
text: Some(text),
..Default::default()
};
let send_result = self.client.channel_messaging().send(&input, opt).await;
match send_result {
Ok(sr) if !sr.message_id.is_empty() => {
self.streams.lock().await.insert(chat_id.clone(), sr.message_id);
}
Ok(_) => {
tracing::warn!("[feishu] 发送流式首帧未返回 message_id");
}
Err(e) => {
tracing::error!("[feishu] 发送流式首帧失败: {e}");
}
}
}
if done {
self.streams.lock().await.remove(&chat_id);
}
Ok(())
}
OutboundAction::ApprovalCard { chat_id, approval_id, title, detail } => {
let card = build_approval_card(&approval_id, &title, &detail);
let document = CardDocument::new(card)
.map_err(|e| anyhow::anyhow!("审批卡片校验失败: {e}"))?;
// encoded_content 是 SDK pub(crate);直接序列化 card(校验已过)
let content = serde_json::to_string(document.card())
.map_err(|e| anyhow::anyhow!("卡片序列化失败: {e}"))?;
let input = larksuite_oapi_sdk_rs::channel::SendInput {
receive_id: Some(chat_id),
card: Some(content),
..Default::default()
};
self.client
.channel_messaging()
.send(&input, opt)
.await
.map_err(|e| anyhow::anyhow!("审批卡片发送失败: {e}"))?;
Ok(())
}
OutboundAction::Notice { chat_id, text } => {
let input = larksuite_oapi_sdk_rs::channel::SendInput {
receive_id: Some(chat_id),
text: Some(text),
..Default::default()
};
self.client
.channel_messaging()
.send(&input, opt)
.await
.map_err(|e| anyhow::anyhow!("通知发送失败: {e}"))?;
Ok(())
}
}
}
}
/// 入站消息:群聊需 @机器人(mentioned_bot)或单聊直发
async fn handle_inbound_message(events: Sender<InboundEvent>, msg: larksuite_oapi_sdk_rs::channel::NormalizedMessage) {
// 只处理文本(text 优先,回退 content)
if msg.message_type != "text" {
tracing::debug!("[feishu] 忽略非文本消息 type={}", msg.message_type);
return;
}
// 群聊必须 @机器人;单聊(p2p)直发
if msg.chat_type == "group" && !msg.mentioned_bot {
tracing::debug!("[feishu] 群聊未 @机器人,忽略");
return;
}
let text = msg.text.clone().unwrap_or_default().trim().to_string();
if text.is_empty() {
return;
}
// 指令:/new 重置会话
if text == "/new" {
let _ = events
.send(InboundEvent::NewSession {
platform: "feishu".into(),
chat_id: msg.chat_id.clone(),
})
.await;
return;
}
let user_id = msg
.sender
.user_id
.as_ref()
.and_then(|u| u.open_id.clone())
.unwrap_or_else(|| "unknown".into());
let _ = events
.send(InboundEvent::Message {
platform: "feishu".into(),
chat_id: msg.chat_id,
user_id,
text,
})
.await;
}
/// 卡片按钮回调:action_value 携带 approval_id + approved
async fn handle_card_action(events: Sender<InboundEvent>, action: larksuite_oapi_sdk_rs::channel::NormalizedCardAction) {
// 按钮约定的 value 形态:{"approval_id": "...", "approved": true|false}
let approval_id = action
.action_value
.get("approval_id")
.and_then(|v| v.as_str())
.unwrap_or_default()
.to_string();
let approved = action
.action_value
.get("approved")
.and_then(|v| v.as_bool())
.unwrap_or(false);
if approval_id.is_empty() {
tracing::warn!("[feishu] 卡片回调缺 approval_id,忽略");
return;
}
let chat_id = action.open_chat_id.clone().unwrap_or_default();
let _ = events
.send(InboundEvent::Approval {
platform: "feishu".into(),
chat_id,
approval_id,
approved,
})
.await;
}
/// 构造审批卡片(Card JSON 1.0):标题 + 详情 + 同意/拒绝按钮
fn build_approval_card(approval_id: &str, title: &str, detail: &str) -> Card {
let detail_md = format!("{}\n\n**审批 ID**: `{approval_id}`", escape_md(detail));
let to_json_value = |v: serde_json::Value| -> larksuite_oapi_sdk_rs::JsonValue {
serde_json::from_value(serde_json::to_value(v).expect("serde_json::Value 可逆"))
.expect("按钮 value 结构兼容")
};
let actions = Action::new()
.action(ActionComponent::Button(
Button::new(larksuite_oapi_sdk_rs::card::v1::Text::plain("✅ 同意"))
.button_type(ButtonType::Primary)
.value(to_json_value(serde_json::json!({"approval_id": approval_id, "approved": true}))),
))
.action(ActionComponent::Button(
Button::new(larksuite_oapi_sdk_rs::card::v1::Text::plain("❌ 拒绝"))
.button_type(ButtonType::Danger)
.value(to_json_value(serde_json::json!({"approval_id": approval_id, "approved": false}))),
));
Card::new()
.header(Header::new(title))
.element(Element::Div(Div::new(
larksuite_oapi_sdk_rs::card::v1::Text::lark_md(detail_md),
)))
.element(Element::Markdown(Markdown::new("---")))
.element(Element::Action(actions))
}
/// Markdown 转义(飞书卡片 lark_md 语法)
fn escape_md(s: &str) -> String {
s.replace('\\', "\\\\")
.replace('*', "\\*")
.replace('_', "\\_")
.replace('`', "\\`")
.replace('[', "\\[")
.replace(']', "\\]")
.replace('<', "\\<")
.replace('>', "\\>")
.replace('(', "\\(")
.replace(')', "\\)")
.replace('#', "\\#")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::FeishuConfig;
/// 构造冒烟:new/offline 均不发网络请求,验证适配器可实例化
#[test]
fn feishu_adapter_constructs() {
let cfg = FeishuConfig {
app_id: "cli_smoke".into(),
app_secret: "secret".into(),
long_connection: true,
};
let adapter = FeishuAdapter::new(&cfg).expect("适配器构造应成功");
assert_eq!(adapter.name(), "feishu");
let offline = FeishuAdapter::offline();
assert_eq!(offline.name(), "feishu");
}
/// 审批卡片构造 + 协议校验(CardDocument::new 是发送前的强校验门槛)
#[test]
fn approval_card_passes_protocol_validation() {
let card = build_approval_card("appr-123", "工具审批请求", "run_command: cargo test");
CardDocument::new(card).expect("审批卡片应通过 Card JSON 1.0 校验");
}
/// 卡片构造包含同意/拒绝按钮及 approval_id 载荷
#[test]
fn approval_card_has_buttons_and_payload() {
let card = build_approval_card("appr-456", "标题", "详情");
let json = serde_json::to_value(&card).unwrap();
let elements = json["elements"].as_array().expect("elements 数组");
// 最后一个元素是 Action 行,内含 2 个按钮
let action = elements.last().unwrap();
let buttons = action["actions"].as_array().expect("actions 数组");
assert_eq!(buttons.len(), 2);
let values: Vec<String> = buttons
.iter()
.map(|b| b["value"]["approval_id"].as_str().unwrap_or_default().to_string())
.collect();
assert!(values.iter().all(|v| v == "appr-456"));
let approved_flags: Vec<bool> = buttons
.iter()
.map(|b| b["value"]["approved"].as_bool().unwrap_or(false))
.collect();
assert!(approved_flags.contains(&true) && approved_flags.contains(&false));
}
/// Markdown 转义:卡片详情中的 markdown 语法字符被转义
#[test]
fn escape_md_works() {
assert_eq!(escape_md("a*b_c`d\\e"), "a\\*b\\_c\\`d\\\\e");
// 走查 #7:链接/标题/引用语法字符也须转义,防用户内容注入卡片结构
assert_eq!(escape_md("[x](http://e)#h <b>(y)"), "\\[x\\]\\(http://e\\)\\#h \\<b\\>\\(y\\)");
}
}
+255
View File
@@ -0,0 +1,255 @@
//! df-im-gateway 主入口:装配 config + store + throttle + commander + adapters
//!
//! 骨架版(P0):Web 适配器真实可跑,事件流全链路打通。
//! 用法:
//! ```text
//! df-im-gateway.exe # 默认配置(env)
//! GATEWAY_RELAY_URL=... GATEWAY_RELAY_TOKEN=... \
//! GATEWAY_DEFAULT_CHANNEL=chan-work df-im-gateway.exe
//! ```
mod adapters;
mod commander;
mod config;
mod feishu;
mod store;
mod throttle;
use std::path::PathBuf;
use tokio::sync::mpsc;
use adapters::{InboundEvent, OutboundAction, PlatformAdapter};
use config::GatewayConfig;
use store::SessionStore;
use throttle::DeltaThrottle;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();
let cfg = GatewayConfig::load()?;
if cfg.channels.is_empty() {
anyhow::bail!("未配置任何 channel 路由(GATEWAY_DEFAULT_CHANNEL),无法启动");
}
let data_dir = std::env::var("GATEWAY_DATA_DIR")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("data"));
std::fs::create_dir_all(&data_dir)?;
let store = SessionStore::open(&data_dir.join("sessions.db"))?;
info_banner(&cfg);
// ── 核心事件总线:适配器 → 核心 ─────────────────────────────
let (inbound_tx, mut inbound_rx) = mpsc::channel::<InboundEvent>(256);
// Web 适配器(四端首个,P0)
let web_adapter = adapters::web::WebAdapter::new();
{
let web_run = web_adapter.clone();
let ev_tx = inbound_tx.clone();
tokio::spawn(async move {
if let Err(e) = web_run.run(ev_tx).await {
tracing::error!("[web] 适配器退出: {e}");
}
});
}
// 飞书适配器(P2):FEISHU_APP_ID/SECRET 配置时启用长连接
let feishu_adapter = cfg.feishu.as_ref().map(|f| match feishu::FeishuAdapter::new(f) {
Ok(adapter) => {
let feishu_run = adapter.clone();
let ev_tx = inbound_tx.clone();
tokio::spawn(async move {
if let Err(e) = feishu_run.run(ev_tx).await {
tracing::error!("[feishu] 适配器退出: {e}");
}
});
tracing::info!("[feishu] 适配器已启动(长连接)");
adapter
}
Err(e) => {
tracing::error!("[feishu] 适配器构造失败(跳过): {e}");
feishu::FeishuAdapter::offline()
}
});
// ── Commander 连接(u-relay)────────────────────────────────
let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false);
let default_channel = cfg
.channels
.iter()
.find(|c| c.prefix.is_empty())
.map(|c| c.channel_id.clone())
.unwrap_or_else(|| cfg.channels[0].channel_id.clone());
// 事件流:commander 任务 → 核心循环
let (ev_tx, mut ev_rx) = mpsc::channel::<df_tunnel::events::TunnelEvent>(256);
// 指令流:核心循环 → commander 任务(透传到 WS)
let (cmd_tx, cmd_rx) = mpsc::channel::<df_tunnel::TunnelCommand>(256);
{
let relay_url = cfg.relay_url.clone();
let token = cfg.relay_token.clone();
let channel = default_channel.clone();
tokio::spawn(async move {
if let Err(e) = commander::run_commander(relay_url, channel, token, ev_tx, cmd_rx, shutdown_rx).await {
tracing::error!("[commander] 连接循环退出: {e}");
}
});
}
// ── 核心分发循环(单 select,持有全部状态,无跨线程共享)─────
let mut throttle = DeltaThrottle::new(cfg.throttle_ms);
// 审批幂等表:approval_id → 已处理标记(P0 必验项 ②:防 pending→approved 二次点击重复转发)
let mut handled_approvals: std::collections::HashSet<String> = std::collections::HashSet::new();
// 平台注册表(走查 #4):显式 platform → 下标映射,替代硬编码 "feishu"=>1 与
// adapters vec 构建顺序的隐式耦合。新增平台(企微/钉钉)只在此处注册一次。
let mut adapters: Vec<Box<dyn adapters::PlatformAdapter>> = vec![Box::new(web_adapter.clone())];
let mut platform_index: std::collections::HashMap<String, usize> =
std::collections::HashMap::from([("web".to_string(), 0usize)]);
if let Some(fs) = feishu_adapter {
platform_index.insert("feishu".to_string(), adapters.len());
adapters.push(Box::new(fs));
}
let platform_of = |platform_or_conv: &str| -> usize {
// conv_id 形如 "platform:chat_id",前缀即平台;未知平台回退 web(0)
let platform = platform_or_conv.split(':').next().unwrap_or("web");
*platform_index.get(platform).unwrap_or(&0)
};
// 停止信号
let shutdown_ctrl = tokio::signal::ctrl_c();
tokio::pin!(shutdown_ctrl);
let cfg_loop = cfg.clone();
loop {
tokio::select! {
_ = &mut shutdown_ctrl => {
tracing::info!("收到 Ctrl+C,优雅退出");
let _ = shutdown_tx.send(true);
break;
}
// ── TunnelEvent(桌面端 → 网关)──
ev = ev_rx.recv() => {
let Some(ev) = ev else { break };
match ev {
df_tunnel::TunnelEvent::TextDelta { conv_id, delta } => {
// 出站按 conv_id 前缀(platform:chat_id)路由回原平台
let flush = throttle.feed(&conv_id, &delta);
if let Some(text) = flush {
let idx = platform_of(&conv_id);
let _ = adapters[idx]
.dispatch(OutboundAction::StreamText { chat_id: conv_id, text, done: false })
.await;
}
}
df_tunnel::TunnelEvent::Completed { conv_id, .. } => {
let tail = throttle.finish(&conv_id);
let full = tail.unwrap_or_default();
let idx = platform_of(&conv_id);
let _ = adapters[idx]
.dispatch(OutboundAction::StreamText { chat_id: conv_id, text: full, done: true })
.await;
}
df_tunnel::TunnelEvent::Approval { conv_id, approval_id, description } => {
let idx = platform_of(&conv_id);
let _ = adapters[idx]
.dispatch(OutboundAction::ApprovalCard {
chat_id: conv_id,
approval_id,
title: "工具审批请求".into(),
detail: description,
})
.await;
}
df_tunnel::TunnelEvent::ToolCall { .. } => {
// P0 骨架不转发 ToolCall 到 IM(卡片噪音),仅日志
tracing::debug!("[core] ToolCall 事件(骨架忽略)");
}
}
}
// ── InboundEvent(IM → 网关)──
inbound = inbound_rx.recv() => {
let Some(ev) = inbound else { break };
match ev {
InboundEvent::Message { platform, chat_id, user_id, text } => {
// 回复适配器:按来源平台回路由
let reply_idx = platform_of(&platform);
if !cfg_loop.whitelist.is_empty() && !cfg_loop.whitelist.contains(&user_id) {
let _ = adapters[reply_idx]
.dispatch(OutboundAction::Notice { chat_id, text: "⛔ 你不在白名单内".into() })
.await;
continue;
}
let Some(channel_id) = cfg_loop.route(&text) else {
let _ = adapters[reply_idx]
.dispatch(OutboundAction::Notice { chat_id, text: "⚠️ 无可用 channel".into() })
.await;
continue;
};
// 会话映射:无则新建(conv_id = platform:chat_id,Web 端 chat_id 同源)
let conv_id = match store.get_conv(&platform, &chat_id) {
Some(m) => m.conv_id,
None => {
let conv = format!("{platform}:{chat_id}");
let _ = store.set_conv(&platform, &chat_id, &conv, channel_id);
// 走查 #3:新建会话时清理流式残留(上轮异常中止无
// Completed 事件,streams entry 永驻会错覆写旧消息)
adapters[reply_idx].reset_stream(&chat_id).await;
conv
}
};
let _ = cmd_tx
.send(df_tunnel::TunnelCommand::Send { conv_id, content: text })
.await;
}
InboundEvent::NewSession { platform, chat_id } => {
let _ = store.remove(&platform, &chat_id);
let reply_idx = platform_of(&platform);
let _ = adapters[reply_idx]
.dispatch(OutboundAction::Notice { chat_id, text: "🔄 已重置会话".into() })
.await;
}
InboundEvent::Approval { platform, chat_id, approval_id, approved } => {
// 幂等:同 approval_id 只处理首次(P0 必验项 ②)
if handled_approvals.contains(&approval_id) {
tracing::debug!("[core] 审批 {approval_id} 已处理,忽略重复点击");
continue;
}
handled_approvals.insert(approval_id.clone());
let conv_id = store
.get_conv(&platform, &chat_id)
.map(|m| m.conv_id)
.unwrap_or_else(|| format!("{platform}:{chat_id}"));
let _ = cmd_tx
.send(df_tunnel::TunnelCommand::Approve { conv_id, approval_id, accepted: approved })
.await;
}
}
}
}
}
Ok(())
}
fn info_banner(cfg: &GatewayConfig) {
println!("╔══════════════════════════════════════════╗");
println!("║ df-im-gateway v0.1.0 (P0 骨架) ║");
println!("╚══════════════════════════════════════════╝");
println!(" relay : {}", cfg.relay_url);
println!(" channels : {} 条路由", cfg.channels.len());
println!(" throttle : {}ms", cfg.throttle_ms);
println!(" whitelist : {} 用户", cfg.whitelist.len());
if cfg.whitelist.is_empty() {
println!(" ⚠️ 白名单为空:任何人可用本网关(开发默认)。生产部署请设置 GATEWAY_WHITELIST!");
}
println!(" feishu : {}", if cfg.feishu.is_some() { "已配置" } else { "未配置(跳过)" });
println!(" web ws : 0.0.0.0:9191");
}
+151
View File
@@ -0,0 +1,151 @@
//! 会话映射存储(IM chat_id ↔ DevFlow conv_id,SQLite + WAL)
//!
//! 建表即带 owner_user_id/platform 列(P0 预留项 ⑤,当前恒默认值,
//! S2 多人场景免迁表)。
//!
//! 表结构:
//! ```sql
//! CREATE TABLE IF NOT EXISTS session_map (
//! platform TEXT NOT NULL, -- feishu/wecom/dingtalk/web
//! chat_id TEXT NOT NULL, -- IM 侧会话(单聊 user_id / 群聊 chat_id)
//! conv_id TEXT NOT NULL, -- DevFlow 会话 ID
//! owner_user_id TEXT NOT NULL DEFAULT '',-- 预留 S2 多人归属(当前恒 '')
//! channel_id TEXT NOT NULL DEFAULT '',-- 路由到的桌面端 channel(预留项 ⑥)
//! created_at INTEGER NOT NULL,
//! updated_at INTEGER NOT NULL,
//! PRIMARY KEY (platform, chat_id)
//! );
//! ```
use std::path::Path;
use std::sync::Arc;
use rusqlite::Connection;
/// 单条会话映射(字段预留 S2 多人/多桌面端,当前核心只读 conv_id)
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct SessionMap {
pub platform: String,
pub chat_id: String,
pub conv_id: String,
pub owner_user_id: String,
pub channel_id: String,
}
#[derive(Clone)]
pub struct SessionStore {
conn: Arc<std::sync::Mutex<Connection>>,
}
impl SessionStore {
/// 打开(或创建)库文件并确保表存在。WAL 模式(与桌面端一致)。
pub fn open(path: &Path) -> anyhow::Result<Self> {
let conn = Connection::open(path)?;
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS session_map (
platform TEXT NOT NULL,
chat_id TEXT NOT NULL,
conv_id TEXT NOT NULL,
owner_user_id TEXT NOT NULL DEFAULT '',
channel_id TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
PRIMARY KEY (platform, chat_id)
);",
)?;
Ok(Self { conn: Arc::new(std::sync::Mutex::new(conn)) })
}
/// 内存库(测试/演示)
#[allow(dead_code)]
pub fn open_in_memory() -> anyhow::Result<Self> {
let conn = Connection::open_in_memory()?;
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS session_map (
platform TEXT NOT NULL, chat_id TEXT NOT NULL, conv_id TEXT NOT NULL,
owner_user_id TEXT NOT NULL DEFAULT '', channel_id TEXT NOT NULL DEFAULT '',
created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL,
PRIMARY KEY (platform, chat_id)
);",
)?;
Ok(Self { conn: Arc::new(std::sync::Mutex::new(conn)) })
}
/// 查 IM 会话映射的 conv_id(未映射返回 None)
pub fn get_conv(&self, platform: &str, chat_id: &str) -> Option<SessionMap> {
let conn = self.conn.lock().ok()?;
conn.query_row(
"SELECT platform, chat_id, conv_id, owner_user_id, channel_id FROM session_map \
WHERE platform = ?1 AND chat_id = ?2",
[platform, chat_id],
|row| {
Ok(SessionMap {
platform: row.get(0)?,
chat_id: row.get(1)?,
conv_id: row.get(2)?,
owner_user_id: row.get(3)?,
channel_id: row.get(4)?,
})
},
)
.ok()
}
/// 写/更新映射(upsert)
pub fn set_conv(&self, platform: &str, chat_id: &str, conv_id: &str, channel_id: &str) -> anyhow::Result<()> {
let conn = self.conn.lock().map_err(|e| anyhow::anyhow!("store 锁中毒: {e}"))?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as i64)
.unwrap_or(0);
conn.execute(
"INSERT INTO session_map (platform, chat_id, conv_id, owner_user_id, channel_id, created_at, updated_at) \
VALUES (?1, ?2, ?3, '', ?4, ?5, ?5) \
ON CONFLICT(platform, chat_id) DO UPDATE SET conv_id = ?3, channel_id = ?4, updated_at = ?5",
rusqlite::params![platform, chat_id, conv_id, channel_id, now],
)?;
Ok(())
}
/// 删除映射(桌面端会话已删,悬空清理)
pub fn remove(&self, platform: &str, chat_id: &str) -> anyhow::Result<()> {
let conn = self.conn.lock().map_err(|e| anyhow::anyhow!("store 锁中毒: {e}"))?;
conn.execute("DELETE FROM session_map WHERE platform = ?1 AND chat_id = ?2", [platform, chat_id])?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn upsert_and_get() {
let store = SessionStore::open_in_memory().unwrap();
store.set_conv("feishu", "oc_123", "conv_abc", "chan-work").unwrap();
let m = store.get_conv("feishu", "oc_123").unwrap();
assert_eq!(m.conv_id, "conv_abc");
assert_eq!(m.channel_id, "chan-work");
assert_eq!(m.owner_user_id, "");
// upsert 覆盖
store.set_conv("feishu", "oc_123", "conv_new", "chan-home").unwrap();
let m = store.get_conv("feishu", "oc_123").unwrap();
assert_eq!(m.conv_id, "conv_new");
// 删除后悬空
store.remove("feishu", "oc_123").unwrap();
assert!(store.get_conv("feishu", "oc_123").is_none());
}
#[test]
fn platform_isolated() {
let store = SessionStore::open_in_memory().unwrap();
store.set_conv("feishu", "c1", "conv1", "ch").unwrap();
store.set_conv("wecom", "c1", "conv2", "ch").unwrap();
assert_eq!(store.get_conv("feishu", "c1").unwrap().conv_id, "conv1");
assert_eq!(store.get_conv("wecom", "c1").unwrap().conv_id, "conv2");
}
}
+103
View File
@@ -0,0 +1,103 @@
//! TextDelta 节流聚合器(IM 频控防护)
//!
//! IM 平台消息 API 频控(飞书 ~50 QPS 应用级、企微 200 次/分钟),逐 delta 刷卡片必被限流。
//! 策略:per-conv 聚合缓冲 + 时间窗(默认 800ms)批量 flush;Completed 时强制冲刷尾包。
//!
//! 设计要点:
//! - 每个 key 一个缓冲,互不干扰
//! - 纯函数式:feed() 返回 Option<全量文本>(窗口到期才返回),不持有回调/不 await
//! - flush 语义:全量文本(IM 卡片更新是「全量内容」语义,非增量)
//! - finish() 冲刷尾包并清缓冲;abort() 丢弃缓冲(错误/停止,不输出残文)
use std::collections::HashMap;
use std::time::{Duration, Instant};
/// 单 key 缓冲
struct ConvBuffer {
/// 累积全量文本
full_text: String,
/// 上次 flush 时刻
last_flush: Instant,
}
/// 节流聚合器
pub struct DeltaThrottle {
window: Duration,
buffers: HashMap<String, ConvBuffer>,
}
impl DeltaThrottle {
pub fn new(window_ms: u64) -> Self {
Self { window: Duration::from_millis(window_ms), buffers: HashMap::new() }
}
/// 喂入 delta。窗口到期则返回 Some(全量文本)(调用方负责外发),否则 None。
pub fn feed(&mut self, key: &str, delta: &str) -> Option<String> {
let now = Instant::now();
let buf = self.buffers.entry(key.to_string()).or_insert_with(|| ConvBuffer {
full_text: String::new(),
last_flush: now,
});
buf.full_text.push_str(delta);
if buf.last_flush.elapsed() >= self.window {
buf.last_flush = now;
Some(buf.full_text.clone())
} else {
None
}
}
/// 会话完成:强制冲刷尾包 + 清缓冲。返回 Some(全量文本) 或(空会话)None。
pub fn finish(&mut self, key: &str) -> Option<String> {
self.buffers.remove(key).map(|buf| buf.full_text)
}
/// 丢弃缓冲(错误/停止,不输出残文)
///
/// 注:当前 TunnelEvent 无 Error/Stopped 变体,核心循环暂无调用点;
/// df-tunnel Phase2 扩展后接入。
#[allow(dead_code)]
pub fn abort(&mut self, key: &str) {
self.buffers.remove(key);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn throttle_window_batch() {
let mut t = DeltaThrottle::new(60_000); // 60s 窗口:模拟窗口永不到期
assert_eq!(t.feed("k", "你"), None);
assert_eq!(t.feed("k", "好"), None);
assert_eq!(t.feed("k", "!"), None);
assert_eq!(t.finish("k"), Some("你好!".to_string()), "finish 冲刷全量");
}
#[test]
fn window_elapsed_flushes() {
let mut t = DeltaThrottle::new(0); // 0ms 窗口:每次 feed 都到期
assert_eq!(t.feed("k", "a"), Some("a".to_string()));
assert_eq!(t.feed("k", "b"), Some("ab".to_string()));
// 缓冲仍在 finish 前的最后 flush 版本;finish 拿走剩余
assert_eq!(t.finish("k"), Some("ab".to_string()));
}
#[test]
fn abort_drops() {
let mut t = DeltaThrottle::new(60_000);
t.feed("k", "残文");
t.abort("k");
assert_eq!(t.finish("k"), None, "abort 后缓冲不存在");
}
#[test]
fn keys_isolated() {
let mut t = DeltaThrottle::new(60_000);
t.feed("a", "1");
t.feed("b", "2");
assert_eq!(t.finish("a"), Some("1".to_string()));
assert_eq!(t.finish("b"), Some("2".to_string()));
}
}