feat: iPet → iAs 完整迁移,Rust 版微信 AI 助手

Phase 1  CLI 框架 + 配置系统
  - clap 子命令: login / listen / send / whoami / usage
  - config.json + env var 替换
  - tracing 日志系统
  - state 持久化(auth/runtime 文件存 + PostgreSQL)

Phase 2  微信通道
  - wechat::client — 完整 iLink Bot HTTP API 实现
  - 扫码登录(终端二维码 + 轮询状态)
  - 长轮询 getupdates / 消息收发 / 监听注册

Phase 3  AI 对话(纯文本 + function calling)
  - LlmProvider trait: DeepSeek + LM Studio 实现
  - SSE 流式解析(text / reasoning / tool_calls delta / usage)
  - Conversation: 消息历史 + chat / chat_with_tools 工具循环

Phase 4  PostgreSQL 集成
  - app_state(认证 KV 存储)
  - chat_records(消息收发记录)
  - llm_usage(Token 用量统计缓存命中率)
  - user_memories(长期记忆持久化)
  - pending_approvals(审批确认码)
  - scheduled_tasks(定时任务表)

Phase 5  一切皆 Skill(工具系统)
  - SkillRegistry: 系统 + 用户 skills 双目录合并
  - SKILL.md 解析器 + 子进程执行器(stdin JSON → stdout)
  - 9 个系统 Skills: datetime / weather / search / email /
    shell / schedule / memos / read_memories / read_summaries
  - ApprovalManager: High 风险技能 → 确认码审批(透明模式)
  - High 风险技能:确认码审批(透明模式)

Phase 6  定时任务调度器

上下文管理
  - ChatSession: checkpoint + token budget (28K) + summaries
  - Token 估算器(中英文自适应)
  - 12h 空闲 → trigger_idle_summary(不入会话)
  - Budget 溢出 → trigger_overflow_summary(入会话 + drain 旧消息)
  - Summarizer: LLM 生成自然语言摘要(fallback 简单截断)
  - 长期记忆 / 摘要 通过 read_memories / read_summaries 工具按需读取

工具调用日志 + Token 统计
  - INFO: 工具名 + 参数 + 结果摘要
  - DEBUG: 子进程 exit/stdout/stderr
  - ias usage --since --until --model 查看用量和缓存命中率
This commit is contained in:
2026-06-01 17:21:43 +08:00
parent 3a2d2769b5
commit b9de3665d9
53 changed files with 9874 additions and 2 deletions
+63
View File
@@ -0,0 +1,63 @@
use clap::{Parser, Subcommand};
/// iAs - 微信 AI 智能助手
#[derive(Parser, Debug)]
#[command(name = "ias", version, about = "微信 AI 智能助手 (Rust)")]
pub struct Cli {
#[command(subcommand)]
pub command: Commands,
}
#[derive(Subcommand, Debug)]
pub enum Commands {
/// 扫码登录微信
Login {
/// 登录超时秒数
#[arg(long, default_value = "480")]
timeout: u64,
},
/// 监听消息(长轮询)
Listen {
/// 启用 AI 自动回复
#[arg(long)]
llm: bool,
/// 仅回显消息(不调用 AI
#[arg(long)]
echo: bool,
},
/// 发送文本消息
Send {
/// 目标用户 ID
#[arg(long)]
to: String,
/// 消息内容
#[arg(long)]
text: String,
/// 上下文令牌
#[arg(long)]
context_token: Option<String>,
},
/// 查看当前登录状态
Whoami,
/// 查看 LLM Token 使用统计
Usage {
/// 起始时间(ISO 格式,默认7天前)
#[arg(long)]
since: Option<String>,
/// 结束时间(ISO 格式,默认现在)
#[arg(long)]
until: Option<String>,
/// 模型名称过滤
#[arg(long)]
model: Option<String>,
},
}
+218
View File
@@ -0,0 +1,218 @@
// 允许未来阶段使用的字段
#![allow(dead_code)]
use serde::Deserialize;
use std::path::Path;
/// iAs 全局配置
#[derive(Debug, Clone, Deserialize)]
pub struct AppConfig {
#[serde(default)]
pub llm: LlmConfig,
#[serde(default)]
pub storage: StorageConfig,
}
#[derive(Debug, Clone, Deserialize)]
pub struct LlmConfig {
/// LLM 提供商:deepseek | lmstudio
#[serde(default = "default_llm_provider")]
pub provider: String,
/// DeepSeek 配置
#[serde(default)]
pub deepseek: DeepSeekConfig,
/// LM Studio 配置
#[serde(default)]
pub lmstudio: LmStudioConfig,
}
#[derive(Debug, Clone, Deserialize)]
pub struct DeepSeekConfig {
#[serde(default = "default_deepseek_api_key")]
pub api_key: String,
#[serde(default = "default_deepseek_model")]
pub model: String,
#[serde(default = "default_deepseek_base_url")]
pub base_url: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct LmStudioConfig {
#[serde(default = "default_lmstudio_base_url")]
pub base_url: String,
#[serde(default = "default_lmstudio_model")]
pub model: String,
#[serde(default)]
pub api_key: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct StorageConfig {
/// 状态文件目录
#[serde(default = "default_state_dir")]
pub state_dir: String,
/// PostgreSQL 连接串(可选,默认为文件存储)
#[serde(default)]
pub database_url: String,
}
// ─── 默认值 ───
fn default_llm_provider() -> String {
"deepseek".to_string()
}
fn default_deepseek_api_key() -> String {
env_or_default("DEEPSEEK_API_KEY", "")
}
fn default_deepseek_model() -> String {
env_or_default("DEEPSEEK_MODEL", "deepseek-v4-flash")
}
fn default_deepseek_base_url() -> String {
env_or_default("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")
}
fn default_lmstudio_base_url() -> String {
env_or_default("LM_STUDIO_BASE_URL", "http://localhost:1234/v1")
}
fn default_lmstudio_model() -> String {
env_or_default("LM_STUDIO_MODEL", "local-model")
}
fn default_state_dir() -> String {
".data/weixin-ilink".to_string()
}
fn env_or_default(key: &str, default: &str) -> String {
std::env::var(key).unwrap_or_else(|_| default.to_string())
}
impl Default for LlmConfig {
fn default() -> Self {
Self {
provider: default_llm_provider(),
deepseek: DeepSeekConfig::default(),
lmstudio: LmStudioConfig::default(),
}
}
}
impl Default for DeepSeekConfig {
fn default() -> Self {
Self {
api_key: default_deepseek_api_key(),
model: default_deepseek_model(),
base_url: default_deepseek_base_url(),
}
}
}
impl Default for LmStudioConfig {
fn default() -> Self {
Self {
base_url: default_lmstudio_base_url(),
model: default_lmstudio_model(),
api_key: String::new(),
}
}
}
impl Default for StorageConfig {
fn default() -> Self {
Self {
state_dir: default_state_dir(),
database_url: String::new(),
}
}
}
impl Default for AppConfig {
fn default() -> Self {
Self {
llm: LlmConfig::default(),
storage: StorageConfig::default(),
}
}
}
impl AppConfig {
/// 从文件加载配置,缺失字段回退到环境变量
pub fn from_file(path: impl AsRef<Path>) -> Self {
let path = path.as_ref();
if path.exists() {
match std::fs::read_to_string(path) {
Ok(content) => {
// 先做环境变量替换
let resolved = resolve_env_vars(&content);
match serde_json::from_str(&resolved) {
Ok(cfg) => return cfg,
Err(e) => {
tracing::warn!("解析配置文件失败 ({}): {},使用默认配置", path.display(), e);
}
}
}
Err(e) => {
tracing::warn!("读取配置文件失败 ({}): {},使用默认配置", path.display(), e);
}
}
}
Self::default()
}
}
/// 替换字符串中的 `${VAR_NAME}` 为环境变量值
fn resolve_env_vars(input: &str) -> String {
let mut result = String::with_capacity(input.len());
let mut chars = input.chars().peekable();
while let Some(c) = chars.next() {
if c == '$' && chars.peek() == Some(&'{') {
chars.next(); // 跳过 '{'
let mut var_name = String::new();
while let Some(&ch) = chars.peek() {
if ch == '}' {
chars.next(); // 跳过 '}'
break;
}
var_name.push(ch);
chars.next();
}
let value = std::env::var(&var_name).unwrap_or_default();
result.push_str(&value);
} else {
result.push(c);
}
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_resolve_env_vars() {
// set_var is unsafe in edition 2024; skip env test
let result = resolve_env_vars("prefix_${TEST_VAR}_suffix");
// TEST_VAR might not be set, so result could be "prefix__suffix"
assert_eq!(result, "prefix__suffix");
}
#[test]
fn test_resolve_missing_env_var() {
let result = resolve_env_vars("prefix_${MISSING_VAR}_suffix");
assert_eq!(result, "prefix__suffix");
}
#[test]
fn test_resolve_no_vars() {
let result = resolve_env_vars("plain text");
assert_eq!(result, "plain text");
}
}
+193
View File
@@ -0,0 +1,193 @@
# iAs 上下文管理系统设计 v2
## 核心理念
1. **长期记忆不固定注入** — 通过工具让 LLM 按需读写
2. **摘要不入上下文** — 除非因上下文空间不足被迫压缩
3. **摘要总结点** — 记录压缩位置,之后的消息始终完整保留
4. **12 小时空闲触发** — 超过 12h 自动压缩,但不插入会话
5. **Token Budget 驱动** — 达到上限才压缩,保持消息完整度优先
---
## 状态模型
```
ChatSession
├── checkpoint: usize = 0 ← 摘要总结点(消息列表中的位置)
├── messages: Vec<Message> ← 全部消息(checkpoint 之前的是历史)
├── summaries: Vec<SummaryEntry> ← 所有摘要
│ ├── pos: usize ← 对应消息列表中的位置
│ ├── text: String ← 摘要内容
│ └── reason: "overflow"|"timeout" ← 生成原因
└── last_user_at: Option<DateTime> ← 上一条用户消息时间
```
### 消息列表视图(LLM 实际看到的内容)
```
TOOL: read_memories / write_memory / read_summaries
↑ LLM 按需调用,不自动注入
VIEW: [system_prompt]
[overflow_summary] ← 仅当因 context 满产生时注入
[messages from checkpoint..] ← 始终完整保留
```
---
## 核心类型
```rust
pub struct SummaryEntry {
pub checkpoint: usize, // 对应 messages 中的位置
pub text: String,
pub reason: SummaryReason,
}
pub enum SummaryReason {
Overflow, // 上下文满了自动压缩
Timeout, // 超过 12 小时空闲
}
pub struct ChatSession {
checkpoint: usize,
messages: Vec<Message>,
summaries: Vec<SummaryEntry>,
last_user_at: Option<DateTime<Utc>>,
token_budget: u32, // 默认 32000
}
```
---
## 核心流程
```
用户发消息 "合肥天气"
1. 距离检查
├── last_user_at 超过 12h
│ └── YES → summarize_all(reason=timeout)
│ ├── 生成摘要存入 summaries(不插入会话)
│ └── checkpoint = messages.len()
2. 构建 LLM 请求
├── system_prompt
├── 【如果有 overflow_summary 且 checkpoint > 0】
│ └── { role: "system", content: "历史摘要:..." }
└── messages[checkpoint..]
3. Token 估算
├── 在 budget 内 → 直接发送
└── 超出 budget → summarize_all(reason=overflow)
├── 生成摘要存入 summaries
├── checkpoint = messages.len()
├── 摘要作为 system message 注入本次请求
└── 重新构建请求(现在只有摘要 + 当前消息)
4. 发送给 LLM
├── LLM 回复(可能包含工具调用)
│ ├── read_memories → 返回长期记忆
│ ├── write_memory → 保存到长期记忆
│ └── read_summaries → 返回 summaries 列表
└── LLM 最终回复 → record_turn(user_msg, assistant_msg)
├── messages.push(user_msg, assistant_msg)
└── last_user_at = now
```
---
## 摘要触发条件对比
| 原因 | 触发条件 | 是否插入会话 | 用途 |
|------|---------|------------|------|
| **Overflow** | token 超 budget | ✅ 自动插入 | 保持上下文不超限 |
| **Timeout** | 距上条消息 >12h | ❌ 不插入(工具可读) | 跨天会话,LLM 自行决定是否查看 |
---
## 长期记忆工具(LLM 按需调用)
```json
{
"name": "read_memories",
"description": "读取用户的长期记忆(跨会话持久保存的个人信息)",
"parameters": { "type": "object", "properties": {} }
}
{
"name": "write_memory",
"description": "写入一条用户的长期记忆",
"parameters": {
"type": "object",
"properties": {
"content": { "type": "string", "description": "记忆内容" }
},
"required": ["content"]
}
}
{
"name": "read_summaries",
"description": "读取历史会话摘要(包括空闲超时产生的摘要)",
"parameters": {
"type": "object",
"properties": {}
}
}
```
这些工具注册为系统 skills(只对 `chat_with_tools` 启用,不暴露给用户)。
## 数据库存储
```sql
-- 长期记忆(跨会话持久)
CREATE TABLE user_memories (
user_id TEXT PRIMARY KEY,
content JSONB NOT NULL DEFAULT '[]',
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
```
摘要和消息会话不持久化(会话消息保存到 `chat_records` 表,但摘要只存在于运行时的 `ChatSession`)。
## 与 iPet 的区别
| 特性 | iPet | iAs v2 |
|------|------|--------|
| 长期记忆 | 自动注入每个请求 | 通过工具按需读写 |
| 滚动摘要 | 自动注入每个请求 | 不注入,工具可读 |
| 会话摘要 | 自动注入 | 不注入,overflow 时例外 |
| 摘要触发 | 定期 + 按消息数 | Token budget + 12h 空闲 |
| 历史完整保留 | ❌ 始终压缩 | ✅ checkpoint 后完整保留 |
## Token 估算
```rust
pub fn estimate_tokens(text: &str) -> u32 {
let cjk = text.chars().filter(|c| c >= '\u{4e00}' && c <= '\u{9fff}').count() as u32;
let other = text.chars().count() as u32 - cjk;
cjk * 15 / 10 + other / 4 + 8 // +8 消息 overhead
}
```
## 实施步骤
| 步骤 | 内容 |
|------|------|
| 1 | ChatSession + SummaryEntry 核心类型 |
| 2 | build_context() — 构建 LLM 消息数组 |
| 3 | record_turn() — 记录对话轮次 |
| 4 | summarize_all() — 摘要生成 + checkpoint 更新 |
| 5 | token_estimate() — Token 估算 |
| 6 | 长期记忆工具(read_memories / write_memory|
| 7 | 摘要读取工具(read_summaries|
| 8 | 集成到 Conversation + main.rs |
+266
View File
@@ -0,0 +1,266 @@
use crate::context::types::ChatSession;
use crate::llm::conversation::Summarizer;
use crate::llm::types::Message;
use std::sync::Arc;
use tokio::sync::Mutex;
/// 估算消息的 token 数
pub fn estimate_tokens(msg: &Message) -> u32 {
let text = msg.content.as_str();
estimate_text_tokens(text) + 8
}
/// 估算纯文本的 token 数
pub fn estimate_text_tokens(text: &str) -> u32 {
if text.is_empty() {
return 0;
}
let cjk_count = text.chars().filter(|c| is_cjk(*c)).count() as u32;
let other_count = text.chars().count() as u32 - cjk_count;
// 中文 ≈ 1.5 token/字,英文 ≈ 0.25 token/字符
cjk_count * 15 / 10 + other_count / 4
}
fn is_cjk(c: char) -> bool {
('\u{4e00}'..='\u{9fff}').contains(&c)
|| ('\u{3400}'..='\u{4dbf}').contains(&c)
|| ('\u{f900}'..='\u{faff}').contains(&c)
}
/// 构建给 LLM 的消息数组(带 token budget 管理)
///
/// 返回消息数组和是否需要摘要的信息
pub async fn build_context(
session: &Arc<Mutex<ChatSession>>,
system_prompt: &str,
) -> Vec<Message> {
let mut s = session.lock().await;
// ── 1. 空闲超时检查(消息到达前由调用方检查)──
// 这里只做构建,超时触发在上层
// ── 2. 构建消息 ──
let mut messages = Vec::new();
let mut used = 0u32;
// 系统提示词
let sys_tokens = estimate_text_tokens(system_prompt);
messages.push(Message::system(system_prompt));
used += sys_tokens + 8;
// overflow 摘要(如果有)
if let Some(summary) = s.latest_overflow_summary() {
let text = format!("【历史对话摘要】\n{}", summary.text);
let tokens = estimate_text_tokens(&text);
if used + tokens < s.token_budget {
messages.push(Message::system(text));
used += tokens + 8;
}
}
// 近期消息(从 checkpoint 开始,倒序添加直到接近 budget)
let recent = s.recent_messages().to_vec();
let mut included = Vec::new();
let mut tail_used = 0u32;
for msg in recent.iter().rev() {
let t = estimate_tokens(msg);
if used + tail_used + t > s.token_budget - 500 {
break;
}
tail_used += t;
included.push(msg.clone());
}
included.reverse();
// 保护:确保最近一条用户消息始终在上下文中
if let Some(last_user) = recent.iter().rev().find(|m| m.role == crate::llm::types::Role::User) {
let already_included = included.iter().any(|m| {
m.role == last_user.role && m.content == last_user.content
});
if !already_included {
let t = estimate_tokens(last_user);
if used + tail_used + t <= s.token_budget + 2000 {
included.push(last_user.clone());
}
}
}
messages.extend(included);
messages
}
/// 触发溢出摘要:压缩 checkpoint 到当前位置之间的新消息
pub async fn trigger_overflow_summary(
session: &Arc<Mutex<ChatSession>>,
summarizer: Option<&Summarizer>,
) -> bool {
let mut s = session.lock().await;
if s.messages.is_empty() {
return false;
}
let prev_checkpoint = s.checkpoint;
let end = s.messages.len();
// checkpoint 始终在 user+assistant 对边界,直接用它作为摘要起点
let start = prev_checkpoint;
if start >= end {
return false;
}
let to_summarize: Vec<_> = s.messages[start..end]
.iter()
.take(40) // 最多 40 条
.cloned()
.collect();
if to_summarize.is_empty() {
return false;
}
let summary_text = generate_summary(&to_summarize, summarizer).await;
s.summaries.push(super::types::SummaryEntry {
checkpoint: prev_checkpoint,
text: summary_text,
reason: super::types::SummaryReason::Overflow,
created_at: chrono::Utc::now(),
});
s.checkpoint = end;
// 清理太旧的消息(checkpoint 之前的保留最近 100 条用于调试)
if s.checkpoint > 200 {
let keep = s.checkpoint - 100;
s.messages.drain(0..keep);
// 调整 checkpoint 和 summary checkpoint
let shift = keep;
s.checkpoint -= shift;
for sum in &mut s.summaries {
if sum.checkpoint >= shift {
sum.checkpoint -= shift;
}
}
}
true
}
/// 触发空闲摘要(不注入,保存到 summaries 供 LLM 工具查询)
pub async fn trigger_idle_summary(
session: &Arc<Mutex<ChatSession>>,
summarizer: Option<&Summarizer>,
) {
let mut s = session.lock().await;
if s.messages.is_empty() {
return;
}
let summary_text = generate_summary(&s.messages, summarizer).await;
let cp = s.messages.len();
s.summaries.push(super::types::SummaryEntry {
checkpoint: cp,
text: summary_text,
reason: super::types::SummaryReason::Timeout,
created_at: chrono::Utc::now(),
});
s.checkpoint = s.messages.len();
// 清理旧消息(checkpoint 之前的全部清除)
if s.checkpoint > 0 {
let shift = s.checkpoint;
s.messages.drain(0..shift);
s.checkpoint = 0;
for sum in &mut s.summaries {
if sum.checkpoint >= shift {
sum.checkpoint -= shift;
}
}
}
}
/// 生成摘要:优先使用 LLM,不可用时回退到简单截断
async fn generate_summary(messages: &[Message], summarizer: Option<&Summarizer>) -> String {
if let Some(summarizer) = summarizer {
if messages.len() >= 3 {
let prompt = build_summary_prompt(messages);
match summarizer(prompt).await {
Ok(text) if !text.is_empty() => {
tracing::info!("LLM 摘要: {:.100}...", text);
return text;
}
_ => tracing::warn!("LLM 摘要失败,回退到简单截断"),
}
}
}
summarize_messages(messages)
}
/// 构建给 LLM 的摘要 prompt
fn build_summary_prompt(messages: &[Message]) -> String {
let convo: String = messages
.iter()
.filter(|m| m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant)
.map(|m| {
let role = if m.role == crate::llm::types::Role::User { "用户" } else { "助手" };
let content: String = m.content.chars().take(200).collect();
format!("{}: {}", role, content)
})
.collect::<Vec<_>>()
.join("\n");
format!(
"请将以下对话压缩为简短摘要,保留关键事实、用户偏好、决定和待办事项。只输出摘要,不要其他内容。\n\n对话:\n{}",
convo
)
}
/// 简单的摘要生成(提取最近 N 条消息的关键内容)
fn summarize_messages(messages: &[Message]) -> String {
// 简单实现:提取用户消息的前 80 个字符作为摘要
let lines: Vec<String> = messages
.iter()
.filter(|m| m.role == crate::llm::types::Role::User || m.role == crate::llm::types::Role::Assistant)
.map(|m| {
let role = match m.role {
crate::llm::types::Role::User => "用户",
crate::llm::types::Role::Assistant => "助手",
_ => "",
};
let preview: String = m.content.chars().take(80).collect();
format!("{}: {}", role, preview)
})
.collect();
if lines.is_empty() {
"暂无历史对话".to_string()
} else {
let result = format!("历史对话摘要({} 条消息):\n{}", lines.len(), lines.join("\n"));
if result.chars().count() > 2000 {
format!("{}...(已截断)", result.chars().take(2000).collect::<String>())
} else {
result
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_estimate_cjk() {
let t = estimate_text_tokens("你好世界");
assert!(t >= 5 && t <= 10, "CJK estimate: {}", t);
}
#[test]
fn test_estimate_ascii() {
let t = estimate_text_tokens("hello world");
assert!(t >= 1 && t <= 6, "ASCII estimate: {}", t);
}
}
+6
View File
@@ -0,0 +1,6 @@
pub mod builder;
pub mod tools;
pub mod types;
pub use tools::MemoryStore;
pub use types::ChatSession;
+92
View File
@@ -0,0 +1,92 @@
use sqlx::PgPool;
use std::sync::Arc;
use tokio::sync::Mutex;
/// 长期记忆管理器(内存 + PostgreSQL 双写)
pub struct MemoryStore {
pool: Option<Arc<PgPool>>,
cache: Arc<Mutex<Vec<String>>>,
}
impl MemoryStore {
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
Self {
pool,
cache: Arc::new(Mutex::new(Vec::new())),
}
}
/// 初始化:从数据库加载记忆到缓存
pub async fn load(&self) {
let pool = match self.pool.as_ref() {
Some(p) => p.clone(),
None => return,
};
if let Ok(rows) = sqlx::query_as::<_, (String,)>(
"SELECT content FROM user_memories WHERE user_id = 'default' ORDER BY id",
)
.fetch_all(pool.as_ref())
.await
{
let mut cache = self.cache.lock().await;
*cache = rows.into_iter().map(|(c,)| c).collect();
}
}
/// 读取所有记忆
pub async fn read(&self) -> String {
let mems = self.cache.lock().await;
if mems.is_empty() {
"暂无长期记忆".to_string()
} else {
mems.iter()
.enumerate()
.map(|(i, m)| format!("{}. {}", i + 1, m))
.collect::<Vec<_>>()
.join("\n")
}
}
/// 写入一条记忆
pub async fn write(&self, content: &str) -> String {
// 写入缓存
self.cache.lock().await.push(content.to_string());
// 写入数据库
if let Some(ref pool) = self.pool {
let _ = sqlx::query(
"INSERT INTO user_memories (user_id, content) VALUES ('default', $1)",
)
.bind(content)
.execute(pool.as_ref())
.await;
}
"已添加长期记忆".to_string()
}
}
/// read_summaries 工具:读取历史摘要
pub async fn read_summaries(
session: &Arc<Mutex<super::types::ChatSession>>,
) -> String {
let s = session.lock().await;
if s.summaries.is_empty() {
"暂无历史摘要".to_string()
} else {
s.summaries
.iter()
.map(|sum| {
let reason = match sum.reason {
super::types::SummaryReason::Overflow => "上下文压缩",
super::types::SummaryReason::Timeout => "空闲超时",
};
format!(
"[{}] {} (原因: {})",
sum.created_at.format("%Y-%m-%d %H:%M"),
sum.text,
reason,
)
})
.collect::<Vec<_>>()
.join("\n---\n")
}
}
+111
View File
@@ -0,0 +1,111 @@
use crate::llm::types::Message;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
/// 摘要原因
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SummaryReason {
/// 上下文 token 超 budget 触发
Overflow,
/// 距上条用户消息超过 12 小时触发
Timeout,
}
/// 一条摘要记录
#[derive(Debug, Clone)]
pub struct SummaryEntry {
/// 摘要对应的消息位置(checkpoint)
pub checkpoint: usize,
/// 摘要文本
pub text: String,
/// 生成原因
pub reason: SummaryReason,
/// 生成时间
pub created_at: DateTime<Utc>,
}
/// 聊天会话状态
#[derive(Debug, Clone)]
pub struct ChatSession {
/// 摘要总结点——此位置之前的消息已被压缩
pub checkpoint: usize,
/// 全部消息(checkpoint 之后的保持完整)
pub messages: Vec<Message>,
/// 所有摘要(keyed by checkpoint
pub summaries: Vec<SummaryEntry>,
/// 上一条用户消息的时间
pub last_user_at: Option<DateTime<Utc>>,
/// Token 预算(默认 28000,留 4000 给回复)
pub token_budget: u32,
/// 12 小时空闲阈值(秒)
pub idle_timeout_secs: i64,
}
impl Default for ChatSession {
fn default() -> Self {
Self {
checkpoint: 0,
messages: Vec::new(),
summaries: Vec::new(),
last_user_at: None,
token_budget: 28000,
idle_timeout_secs: 12 * 3600,
}
}
}
impl ChatSession {
pub fn new() -> Self {
Self::default()
}
/// 记录一条用户消息
pub fn add_user(&mut self, content: String) {
self.messages.push(Message::user(content));
self.last_user_at = Some(Utc::now());
}
/// 记录一条助手消息
pub fn add_assistant(&mut self, content: String) {
self.messages.push(Message::assistant(content));
}
/// 记录一条带 tool_calls 的助手消息
pub fn add_assistant_tool_calls(&mut self, tool_calls: Vec<crate::llm::types::ToolCall>) {
self.messages.push(Message::assistant_with_tool_calls(tool_calls));
}
/// 记录一条工具结果
pub fn add_tool_result(&mut self, tool_call_id: &str, name: &str, result: &str) {
self.messages.push(Message::tool_result(tool_call_id, name, result));
}
/// 是否因空闲超过阈值需要摘要
pub fn is_idle_timeout(&self) -> bool {
if let Some(last) = self.last_user_at {
let elapsed = Utc::now().signed_duration_since(last).num_seconds();
elapsed > self.idle_timeout_secs
} else {
false
}
}
/// 获取检查点之后的消息
pub fn recent_messages(&self) -> &[Message] {
if self.checkpoint < self.messages.len() {
&self.messages[self.checkpoint..]
} else {
&[]
}
}
/// 是否有 overflow 摘要可以注入
pub fn has_overflow_summary(&self) -> bool {
self.summaries.iter().any(|s| s.reason == SummaryReason::Overflow)
}
/// 获取最新的 overflow 摘要
pub fn latest_overflow_summary(&self) -> Option<&SummaryEntry> {
self.summaries.iter().rev().find(|s| s.reason == SummaryReason::Overflow)
}
}
+49
View File
@@ -0,0 +1,49 @@
pub mod models;
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use std::time::Duration;
/// 数据库管理器
pub struct Database {
pool: PgPool,
}
impl Database {
/// 从 DATABASE_URL 环境变量连接并运行迁移
pub async fn connect() -> Result<Self, String> {
let database_url =
std::env::var("DATABASE_URL").map_err(|_| "请设置 DATABASE_URL 环境变量".to_string())?;
let pool = PgPoolOptions::new()
.max_connections(5)
.acquire_timeout(Duration::from_secs(10))
.connect(&database_url)
.await
.map_err(|e| format!("连接数据库失败: {}", e))?;
// 运行迁移
sqlx::migrate!("./migrations")
.run(&pool)
.await
.map_err(|e| format!("数据库迁移失败: {}", e))?;
tracing::info!("数据库连接成功,迁移已完成");
Ok(Self { pool })
}
/// 获取连接池引用
pub fn pool(&self) -> &PgPool {
&self.pool
}
/// 检查数据库是否可用
pub async fn health_check(&self) -> Result<(), String> {
sqlx::query("SELECT 1")
.execute(&self.pool)
.await
.map_err(|e| format!("数据库健康检查失败: {}", e))?;
Ok(())
}
}
+193
View File
@@ -0,0 +1,193 @@
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use crate::state::AuthState;
// ─── 认证状态 ───
/// 从数据库加载认证状态
pub async fn load_auth(pool: &PgPool) -> Option<AuthState> {
let row = sqlx::query_as::<_, (serde_json::Value,)>(
"SELECT value FROM app_state WHERE key = 'auth'",
)
.fetch_optional(pool)
.await
.ok()??;
serde_json::from_value(row.0).ok()
}
/// 保存认证状态到数据库
pub async fn save_auth(pool: &PgPool, auth: &AuthState) -> Result<(), String> {
let value = serde_json::to_value(auth).map_err(|e| format!("序列化 auth 失败: {}", e))?;
sqlx::query(
r#"
INSERT INTO app_state (key, value, updated_at)
VALUES ('auth', $1, NOW())
ON CONFLICT (key) DO UPDATE
SET value = EXCLUDED.value, updated_at = NOW()
"#,
)
.bind(&value)
.execute(pool)
.await
.map_err(|e| format!("保存 auth 到数据库失败: {}", e))?;
Ok(())
}
// ─── 聊天记录 ───
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct ChatRecord {
pub id: i64,
pub created_at: DateTime<Utc>,
pub direction: String,
pub user_id: String,
pub account_id: String,
pub text: String,
pub source: String,
pub context_token: String,
pub message_id: String,
}
/// 插入一条聊天记录
pub async fn insert_chat_record(
pool: &PgPool,
direction: &str,
user_id: &str,
account_id: &str,
text: &str,
source: &str,
context_token: Option<&str>,
message_id: &str,
) -> Result<i64, String> {
let ctx = context_token.unwrap_or("");
let row: (i64,) = sqlx::query_as(
r#"
INSERT INTO chat_records (direction, user_id, account_id, text, source, context_token, message_id)
VALUES ($1, $2, $3, $4, $5, $6, $7)
RETURNING id
"#,
)
.bind(direction)
.bind(user_id)
.bind(account_id)
.bind(text)
.bind(source)
.bind(ctx)
.bind(message_id)
.fetch_one(pool)
.await
.map_err(|e| format!("插入聊天记录失败: {}", e))?;
Ok(row.0)
}
/// 查询最近的聊天记录(用户维度)
pub async fn list_recent_chat_records(
pool: &PgPool,
user_id: &str,
limit: i64,
) -> Result<Vec<ChatRecord>, String> {
let records = sqlx::query_as::<_, ChatRecord>(
r#"
SELECT id, created_at, direction, user_id, account_id, text, source, context_token, message_id
FROM chat_records
WHERE user_id = $1
ORDER BY created_at DESC
LIMIT $2
"#,
)
.bind(user_id)
.bind(limit)
.fetch_all(pool)
.await
.map_err(|e| format!("查询聊天记录失败: {}", e))?;
Ok(records)
}
// ─── LLM 用量 ───
/// 插入一条 LLM 用量记录
pub async fn insert_llm_usage(
pool: &PgPool,
user_id: &str,
model: &str,
provider: &str,
prompt_tokens: u32,
completion_tokens: u32,
cache_hit_tokens: u32,
cache_miss_tokens: u32,
) -> Result<(), String> {
let total = prompt_tokens + completion_tokens;
sqlx::query(
r#"
INSERT INTO llm_usage (user_id, model, provider, prompt_tokens, completion_tokens, total_tokens, cache_hit_tokens, cache_miss_tokens)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
"#,
)
.bind(user_id)
.bind(model)
.bind(provider)
.bind(prompt_tokens as i32)
.bind(completion_tokens as i32)
.bind(total as i32)
.bind(cache_hit_tokens as i32)
.bind(cache_miss_tokens as i32)
.execute(pool)
.await
.map_err(|e| format!("插入 LLM 用量失败: {}", e))?;
Ok(())
}
/// LLM 用量聚合统计
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct LlmUsageStats {
pub total_calls: i64,
pub total_prompt_tokens: i64,
pub total_completion_tokens: i64,
pub total_tokens: i64,
pub total_cache_hit: i64,
pub total_cache_miss: i64,
}
/// 查询 LLM 用量统计(按时间范围过滤)
pub async fn query_llm_usage_stats(
pool: &PgPool,
since: Option<chrono::DateTime<chrono::Utc>>,
until: Option<chrono::DateTime<chrono::Utc>>,
model_filter: Option<&str>,
) -> Result<LlmUsageStats, String> {
let since = since.unwrap_or_else(|| {
chrono::Utc::now() - chrono::Duration::days(7)
});
let until = until.unwrap_or_else(chrono::Utc::now);
let row = sqlx::query_as::<_, LlmUsageStats>(
r#"
SELECT
COUNT(*)::bigint as total_calls,
COALESCE(SUM(prompt_tokens), 0)::bigint as total_prompt_tokens,
COALESCE(SUM(completion_tokens), 0)::bigint as total_completion_tokens,
COALESCE(SUM(total_tokens), 0)::bigint as total_tokens,
COALESCE(SUM(cache_hit_tokens), 0)::bigint as total_cache_hit,
COALESCE(SUM(cache_miss_tokens), 0)::bigint as total_cache_miss
FROM llm_usage
WHERE created_at >= $1 AND created_at <= $2
AND ($3::text IS NULL OR model = $3)
"#,
)
.bind(since)
.bind(until)
.bind(model_filter)
.fetch_one(pool)
.await
.map_err(|e| format!("查询 LLM 用量失败: {}", e))?;
Ok(row)
}
+242
View File
@@ -0,0 +1,242 @@
use crate::context::builder;
use crate::context::types::ChatSession;
use crate::llm::provider::{create_provider, BoxedProvider, StreamReceiver};
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Mutex;
pub type ToolExecutor =
Arc<dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync>;
/// 摘要生成器:接收要摘要的消息文本 → 返回 LLM 生成的摘要
pub type Summarizer =
Arc<dyn Fn(String) -> Pin<Box<dyn Future<Output = Result<String, String>> + Send>> + Send + Sync>;
pub struct Conversation {
config: ConversationConfig,
provider: BoxedProvider,
session: Arc<Mutex<ChatSession>>,
tool_executor: Option<ToolExecutor>,
}
impl Conversation {
pub fn new(config: ConversationConfig) -> Result<Self, String> {
let provider = create_provider(&config)?;
Ok(Self {
config,
provider,
session: Arc::new(Mutex::new(ChatSession::new())),
tool_executor: None,
})
}
pub fn with_provider(config: ConversationConfig, provider: BoxedProvider) -> Self {
Self {
config,
provider,
session: Arc::new(Mutex::new(ChatSession::new())),
tool_executor: None,
}
}
pub fn set_tool_executor(&mut self, executor: ToolExecutor) { self.tool_executor = Some(executor); }
pub fn session(&self) -> Arc<Mutex<ChatSession>> { Arc::clone(&self.session) }
pub fn model(&self) -> &str { &self.config.model }
pub fn provider_name(&self) -> &str { self.provider.name() }
pub fn spec(&self) -> &ConversationConfig { &self.config }
pub fn summarizer(&self) -> Summarizer { self.create_summarizer() }
/// 创建 LLM 摘要生成器(用当前 provider 做非流式调用)
pub fn create_summarizer(&self) -> Summarizer {
let provider = self.provider.clone();
Arc::new(move |prompt: String| {
let provider = provider.clone();
Box::pin(async move {
let msgs = vec![Message::system("你是一个对话摘要助手,用中文输出简洁摘要。"), Message::user(prompt)];
let mut rx = provider.chat_stream(&ConversationConfig::default(), &msgs).await
.map_err(|e| format!("摘要请求失败: {}", e))?;
let mut text = String::new();
while let Some(chunk) = rx.recv().await {
match chunk {
StreamChunk::Text(t) => text.push_str(&t),
StreamChunk::Done { text: full, .. } => {
if !full.is_empty() { text = full; }
break;
}
StreamChunk::Error(e) => return Err(e),
_ => {}
}
}
Ok(text)
})
})
}
/// 单次对话(无工具)
pub async fn chat(&self, user_message: String) -> Result<ChatHandle, String> {
// 空闲超时检查
{
let s = self.session.lock().await;
if s.is_idle_timeout() {
drop(s);
builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await;
}
}
self.session.lock().await.add_user(user_message.clone());
let messages = builder::build_context(&self.session, &self.config.system_prompt).await;
let rx = self.provider.chat_stream(&self.config, &messages).await?;
Ok(ChatHandle {
rx: Some(rx), full_text: String::new(), tool_calls: None, usage: None,
session: Arc::clone(&self.session),
})
}
/// 对话 + 工具循环(带上下文管理)
pub async fn chat_with_tools(
&self,
user_message: String,
) -> Result<(String, bool, Option<Usage>), String> {
// ── 空闲超时检查 ──
{
let s = self.session.lock().await;
if s.is_idle_timeout() {
drop(s);
builder::trigger_idle_summary(&self.session, Some(&self.summarizer())).await;
tracing::info!("⏰ 检测到 12h 空闲,已生成摘要");
}
}
self.session.lock().await.add_user(user_message.clone());
let mut used_tools = false;
let mut last_usage = None;
let mut turn_count = 0u32;
loop {
turn_count += 1;
if turn_count > 5 {
return Err("工具调用次数过多(最多5轮)".to_string());
}
let messages = builder::build_context(&self.session, &self.config.system_prompt).await;
let rx = self.provider.chat_stream(&self.config, &messages).await?;
let mut full_text = String::new();
let mut tool_calls: Option<Vec<ToolCall>> = None;
let mut rx = rx;
while let Some(chunk) = rx.recv().await {
match chunk {
StreamChunk::Text(t) => full_text.push_str(&t),
StreamChunk::Done { text, tool_calls: tc, usage, .. } => {
if !text.is_empty() { full_text = text; }
tool_calls = tc;
last_usage = usage;
break;
}
StreamChunk::Error(e) => return Err(e),
_ => {}
}
}
if let Some(calls) = tool_calls {
if !calls.is_empty() {
used_tools = true;
tracing::info!(
"🔧 LLM 请求 {} 个工具: {}",
calls.len(),
calls.iter().map(|c| format!("{}({:.80})", c.function.name, c.function.arguments)).collect::<Vec<_>>().join(", ")
);
self.session.lock().await.add_assistant_tool_calls(calls.clone());
for tc in &calls {
let result = match &self.tool_executor {
Some(exec) => match exec(&tc.function.name, &tc.function.arguments).await {
Ok(r) => r,
Err(e) => format!("执行失败: {}", e),
},
None => "工具执行器未配置".to_string(),
};
tracing::info!("📦 {} → {:.150}", tc.function.name, result);
self.session.lock().await.add_tool_result(&tc.id, &tc.function.name, &result);
}
continue;
}
}
// 无工具调用:记录回复,检查是否需要溢出摘要
if !full_text.is_empty() {
self.session.lock().await.add_assistant(full_text.clone());
}
// 检查 token 是否接近预算,触发溢出摘要
{
let s = self.session.lock().await;
let recent = s.recent_messages();
let estimated: u32 = recent.iter().map(|m| builder::estimate_tokens(m)).sum();
if estimated > s.token_budget {
drop(s);
builder::trigger_overflow_summary(&self.session, Some(&self.summarizer())).await;
tracing::info!("📦 上下文超预算,已触发溢出摘要");
}
}
return Ok((full_text, used_tools, last_usage));
}
}
}
// ─── ChatHandle ───
pub struct ChatHandle {
rx: Option<StreamReceiver>,
full_text: String,
tool_calls: Option<Vec<ToolCall>>,
usage: Option<Usage>,
session: Arc<Mutex<ChatSession>>,
}
impl ChatHandle {
pub async fn next_chunk(&mut self) -> Option<StreamChunk> {
let chunk = self.rx.as_mut()?.recv().await;
match &chunk {
Some(StreamChunk::Text(t)) => self.full_text.push_str(t),
Some(StreamChunk::Done { text, tool_calls, usage, .. }) => {
if !text.is_empty() { self.full_text = text.clone(); }
self.tool_calls = tool_calls.clone();
self.usage = usage.clone();
if !self.full_text.is_empty() {
let mut s = self.session.lock().await;
s.add_assistant(self.full_text.clone());
}
self.rx = None;
}
Some(StreamChunk::Error(_)) => { self.rx = None; }
_ => {}
}
chunk
}
pub async fn consume(mut self) -> Result<String, String> {
while let Some(chunk) = self.next_chunk().await {
if let StreamChunk::Error(e) = chunk { return Err(e); }
}
Ok(self.full_text.clone())
}
pub fn current_text(&self) -> &str { &self.full_text }
pub fn get_tool_calls(&self) -> Option<&Vec<ToolCall>> { self.tool_calls.as_ref() }
pub fn get_usage(&self) -> Option<&Usage> { self.usage.as_ref() }
}
pub const DEFAULT_SYSTEM_PROMPT: &str = "You are a concise and helpful assistant replying in Chinese. \
Keep replies practical and natural. \
Unless the user asks for detail, keep most replies under 120 Chinese characters. \
\
如果需要了解用户的长期记忆或历史对话摘要,使用 read_memories / read_summaries 工具。\
如果用户提到重要的个人信息或约定,使用 write_memory 工具记录下来。";
+195
View File
@@ -0,0 +1,195 @@
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
use crate::llm::types::{ConversationConfig, Message, StreamChunk, ToolCall, Usage};
use async_trait::async_trait;
use reqwest::Client as HttpClient;
use std::collections::BTreeMap;
use tokio::sync::mpsc;
pub struct DeepSeekProvider {
http: HttpClient,
api_key: String,
base_url: String,
}
impl DeepSeekProvider {
pub fn new() -> Result<Self, String> {
let api_key = std::env::var("DEEPSEEK_API_KEY")
.map_err(|_| "请设置 DEEPSEEK_API_KEY 环境变量".to_string())?;
let base_url = std::env::var("DEEPSEEK_BASE_URL")
.unwrap_or_else(|_| "https://api.deepseek.com/v1".to_string());
Ok(Self {
http: HttpClient::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
api_key,
base_url,
})
}
}
#[async_trait]
impl LlmProvider for DeepSeekProvider {
fn name(&self) -> &str { "deepseek" }
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String> {
let url = format!("{}/chat/completions", self.base_url);
let mut body = serde_json::json!({
"model": config.model,
"messages": messages,
"temperature": config.temperature,
"max_tokens": config.max_tokens,
"stream": true,
});
// 如果有工具,添加 tools 参数
if let Some(ref tools) = config.tools {
body["tools"] = serde_json::Value::Array(tools.clone());
}
if config.thinking {
body["thinking"] = serde_json::json!({"type": "enabled"});
body["reasoning_effort"] = serde_json::json!("high");
body.as_object_mut().map(|obj| {
obj.remove("temperature");
obj.remove("top_p");
});
}
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
let http = self.http.clone();
let api_key = self.api_key.clone();
tokio::spawn(async move {
if let Err(e) = stream_deepseek(http, &url, &api_key, body, tx.clone()).await {
let _ = tx.send(StreamChunk::Error(e)).await;
}
});
Ok(rx)
}
}
async fn stream_deepseek(
http: HttpClient,
url: &str,
api_key: &str,
body: serde_json::Value,
tx: StreamSender,
) -> Result<(), String> {
let response = http
.post(url)
.header("Authorization", format!("Bearer {}", api_key))
.header("Content-Type", "application/json")
.header("Accept", "text/event-stream")
.json(&body)
.send()
.await
.map_err(|e| format!("请求失败: {}", e))?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
return Err(format!("HTTP {}: {}", status, text));
}
let mut full_text = String::new();
let mut full_reasoning = String::new();
let mut usage: Option<Usage> = None;
let mut tool_call_builders: BTreeMap<i32, (Option<String>, Option<String>, String)> = BTreeMap::new();
use futures_util::StreamExt;
let mut buf = String::new();
let mut stream = response.bytes_stream();
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
let text = String::from_utf8_lossy(&chunk);
buf.push_str(&text);
while let Some(line_end) = buf.find('\n') {
let line = buf[..line_end].to_string();
buf = buf[line_end + 1..].to_string();
let trimmed = line.trim();
if trimmed.is_empty() { continue; }
if let Some(parsed) = parse_chat_chunk(trimmed) {
match parsed {
ParsedChunk::Text(t) => {
full_text.push_str(&t);
let _ = tx.send(StreamChunk::Text(t)).await;
}
ParsedChunk::Reasoning(r) => {
full_reasoning.push_str(&r);
let _ = tx.send(StreamChunk::Reasoning(r)).await;
}
ParsedChunk::ToolCallDelta { index, id, name, arguments } => {
let entry = tool_call_builders.entry(index).or_insert((None, None, String::new()));
if let Some(call_id) = id { entry.0 = Some(call_id); }
if let Some(call_name) = name { entry.1 = Some(call_name); }
entry.2.push_str(&arguments);
}
ParsedChunk::FinishReason(_) => {}
ParsedChunk::Usage(u) => { usage = Some(u); }
}
}
}
}
// 检查最后的 buffer
if !buf.trim().is_empty() {
if let Some(parsed) = parse_chat_chunk(buf.trim()) {
match parsed {
ParsedChunk::Text(t) => full_text.push_str(&t),
ParsedChunk::ToolCallDelta { index, id, name, arguments } => {
let entry = tool_call_builders.entry(index).or_insert((None, None, String::new()));
if let Some(call_id) = id { entry.0 = Some(call_id); }
if let Some(call_name) = name { entry.1 = Some(call_name); }
entry.2.push_str(&arguments);
}
ParsedChunk::Usage(u) => { usage = Some(u); }
_ => {}
}
}
}
// 构建 tool_calls
let tool_calls: Option<Vec<ToolCall>> = if tool_call_builders.is_empty() {
None
} else {
Some(
tool_call_builders
.into_values()
.filter_map(|(id, name, args)| {
let name = name?;
let id = id.unwrap_or_else(|| format!("call_{}", rand::random::<u32>()));
Some(ToolCall {
id,
call_type: "function".to_string(),
function: crate::llm::types::ToolFunctionCall {
name,
arguments: args,
},
})
})
.collect(),
)
};
let _ = tx
.send(StreamChunk::Done {
text: full_text,
reasoning: if full_reasoning.is_empty() { None } else { Some(full_reasoning) },
tool_calls,
usage,
})
.await;
Ok(())
}
+157
View File
@@ -0,0 +1,157 @@
use crate::llm::provider::{parse_chat_chunk, LlmProvider, ParsedChunk, StreamReceiver, StreamSender};
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
use async_trait::async_trait;
use reqwest::Client as HttpClient;
use tokio::sync::mpsc;
/// LM Studio 提供商(OpenAI 兼容接口)
pub struct LmStudioProvider {
http: HttpClient,
base_url: String,
api_key: String,
}
impl LmStudioProvider {
pub fn new() -> Result<Self, String> {
let base_url = std::env::var("LM_STUDIO_BASE_URL")
.unwrap_or_else(|_| "http://localhost:1234/v1".to_string());
// 标准化 base_url: 确保有协议和 /v1 路径
let base_url = normalize_base_url(&base_url);
let api_key = std::env::var("LM_STUDIO_API_KEY").unwrap_or_default();
Ok(Self {
http: HttpClient::builder()
.timeout(std::time::Duration::from_secs(120))
.build()
.map_err(|e| format!("创建 HTTP client 失败: {}", e))?,
base_url,
api_key,
})
}
}
#[async_trait]
impl LlmProvider for LmStudioProvider {
fn name(&self) -> &str {
"lmstudio"
}
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String> {
let url = format!("{}/chat/completions", self.base_url);
let body = serde_json::json!({
"model": config.model,
"messages": messages,
"temperature": config.temperature,
"max_tokens": config.max_tokens,
"stream": true,
});
let (tx, rx) = mpsc::channel::<StreamChunk>(64);
let http = self.http.clone();
let api_key = self.api_key.clone();
tokio::spawn(async move {
if let Err(e) = stream_openai(http, &url, &api_key, body, tx.clone()).await {
let _ = tx.send(StreamChunk::Error(e)).await;
}
});
Ok(rx)
}
}
/// 与 DeepSeek 公用相同的 SSE 流式处理,但去除 thinking 相关字段
async fn stream_openai(
http: HttpClient,
url: &str,
api_key: &str,
body: serde_json::Value,
tx: StreamSender,
) -> Result<(), String> {
let mut req = http.post(url).header("Content-Type", "application/json");
if !api_key.is_empty() {
req = req.header("Authorization", format!("Bearer {}", api_key));
}
let response = req
.json(&body)
.send()
.await
.map_err(|e| format!("请求失败: {}", e))?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
let _ = tx.send(StreamChunk::Error(format!("HTTP {}: {}", status, text))).await;
return Err(format!("HTTP {}: {}", status, text));
}
let mut full_text = String::new();
let mut buf = String::new();
let mut stream = response.bytes_stream();
use futures_util::StreamExt;
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(|e| format!("读取流失败: {}", e))?;
let text = String::from_utf8_lossy(&chunk);
buf.push_str(&text);
while let Some(line_end) = buf.find('\n') {
let line = buf[..line_end].to_string();
buf = buf[line_end + 1..].to_string();
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
// LM Studio 使用标准 OpenAI SSE 格式
if let Some(parsed) = parse_chat_chunk(trimmed) {
match parsed {
ParsedChunk::Text(t) => {
full_text.push_str(&t);
let _ = tx.send(StreamChunk::Text(t)).await;
}
ParsedChunk::Reasoning(r) => {
let _ = tx.send(StreamChunk::Reasoning(r)).await;
}
_ => {}
}
}
}
}
let _ = tx
.send(StreamChunk::Done {
text: full_text,
reasoning: None,
tool_calls: None,
usage: None,
})
.await;
Ok(())
}
/// 确保 base_url 有协议前缀和 /v1 路径
fn normalize_base_url(raw: &str) -> String {
let raw = raw.trim();
let with_protocol = if raw.starts_with("http://") || raw.starts_with("https://") {
raw.to_string()
} else {
format!("http://{}", raw)
};
// 去掉末尾 /
let trimmed = with_protocol.trim_end_matches('/');
// 确保路径以 /v1 结尾
if trimmed.ends_with("/v1") {
trimmed.to_string()
} else {
format!("{}/v1", trimmed)
}
}
+8
View File
@@ -0,0 +1,8 @@
pub mod conversation;
pub mod deepseek;
pub mod lmstudio;
pub mod provider;
pub mod types;
pub use conversation::{Conversation, Summarizer, ToolExecutor, DEFAULT_SYSTEM_PROMPT};
pub use types::{ConversationConfig, Message, Role, StreamChunk, Usage};
+187
View File
@@ -0,0 +1,187 @@
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
use serde::Deserialize;
use async_trait::async_trait;
use std::sync::Arc;
use tokio::sync::mpsc;
/// 流式返回:每个 StreamChunk 是 delta 或控制信号
pub type StreamReceiver = mpsc::Receiver<StreamChunk>;
pub type StreamSender = mpsc::Sender<StreamChunk>;
/// LLM 提供商抽象
#[async_trait]
pub trait LlmProvider: Send + Sync {
/// 提供商名称
fn name(&self) -> &str;
/// 发起流式对话
async fn chat_stream(
&self,
config: &ConversationConfig,
messages: &[Message],
) -> Result<StreamReceiver, String>;
}
// ─── 内置提供商注册 ───
pub type BoxedProvider = Arc<dyn LlmProvider>;
/// 从配置创建恰当的提供商
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
// 根据 model 前缀或环境变量选择
let provider = std::env::var("LLM_PROVIDER")
.unwrap_or_else(|_| "deepseek".to_string())
.to_lowercase();
match provider.as_str() {
"deepseek" => Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?)),
"lmstudio" => Ok(Arc::new(super::lmstudio::LmStudioProvider::new()?)),
_ => Err(format!("不支持的 LLM 提供商: {}", provider)),
}
}
// ─── 内部:SSE 解析工具 ───
/// 从字节流中解析 SSE 行
pub(crate) fn parse_sse_line(line: &str) -> Option<(String, String)> {
// 支持两种格式:
// data: {...}
// event: ...
let line = line.trim();
if line.is_empty() || line.starts_with(':') {
return None;
}
if let Some(pos) = line.find(": ") {
let field = line[..pos].trim().to_string();
let value = line[pos + 2..].trim_start().to_string();
Some((field, value))
} else if let Some(pos) = line.find(':') {
let field = line[..pos].trim().to_string();
let value = line[pos + 1..].trim_start().to_string();
Some((field, value))
} else {
None
}
}
/// 流式块解析结果
pub(crate) enum ParsedChunk {
Text(String),
Reasoning(String),
ToolCallDelta {
index: i32,
id: Option<String>,
name: Option<String>,
arguments: String,
},
FinishReason(String),
Usage(super::types::Usage),
}
/// 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta
pub(crate) fn parse_chat_chunk(
line: &str,
) -> Option<ParsedChunk> {
if !line.starts_with("data: ") {
return None;
}
let data = line["data: ".len()..].trim();
if data == "[DONE]" {
return None;
}
#[derive(Deserialize)]
struct ToolCallDelta {
#[serde(default)]
index: Option<i32>,
#[serde(default)]
id: Option<String>,
#[serde(rename = "type", default)]
call_type: Option<String>,
#[serde(default)]
function: Option<ToolFunctionDelta>,
}
#[derive(Deserialize)]
struct ToolFunctionDelta {
#[serde(default)]
name: Option<String>,
#[serde(default)]
arguments: Option<String>,
}
#[derive(Deserialize)]
struct Delta {
#[serde(default)]
content: Option<String>,
#[serde(default)]
reasoning_content: Option<String>,
#[serde(default)]
tool_calls: Option<Vec<ToolCallDelta>>,
}
#[derive(Deserialize)]
struct ChunkChoice {
delta: Delta,
#[serde(default)]
finish_reason: Option<String>,
}
#[derive(Deserialize)]
struct ChunkResponse {
choices: Vec<ChunkChoice>,
#[serde(default)]
usage: Option<super::types::Usage>,
}
let parsed: ChunkResponse = match serde_json::from_str(data) {
Ok(p) => p,
Err(_) => return None,
};
// 提取 usage(可能在最后一个 chunk
if let Some(ref usage) = parsed.usage {
if usage.total_tokens > 0 {
return Some(ParsedChunk::Usage(usage.clone()));
}
}
for choice in parsed.choices {
// 工具调用 delta
if let Some(tool_calls) = &choice.delta.tool_calls {
for tc in tool_calls {
let idx = tc.index.unwrap_or(0);
let args = tc.function.as_ref()
.and_then(|f| f.arguments.clone())
.unwrap_or_default();
let name = tc.function.as_ref().and_then(|f| f.name.clone());
return Some(ParsedChunk::ToolCallDelta {
index: idx,
id: tc.id.clone(),
name,
arguments: args,
});
}
}
if let Some(reasoning) = &choice.delta.reasoning_content {
if !reasoning.is_empty() {
return Some(ParsedChunk::Reasoning(reasoning.clone()));
}
}
if let Some(content) = &choice.delta.content {
if !content.is_empty() {
return Some(ParsedChunk::Text(content.clone()));
}
}
if let Some(reason) = &choice.finish_reason {
if !reason.is_empty() {
return Some(ParsedChunk::FinishReason(reason.clone()));
}
}
}
None
}
+125
View File
@@ -0,0 +1,125 @@
use serde::{Deserialize, Serialize};
// ─── 角色 ───
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
User,
Assistant,
Tool,
}
// ─── 消息 ───
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
#[serde(default)]
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self { role: Role::System, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
}
pub fn user(content: impl Into<String>) -> Self {
Self { role: Role::User, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
}
pub fn assistant(content: impl Into<String>) -> Self {
Self { role: Role::Assistant, content: content.into(), tool_call_id: None, name: None, tool_calls: None }
}
pub fn assistant_with_tool_calls(tool_calls: Vec<ToolCall>) -> Self {
Self { role: Role::Assistant, content: String::new(), tool_call_id: None, name: None, tool_calls: Some(tool_calls) }
}
pub fn tool_result(tool_call_id: &str, name: &str, result: &str) -> Self {
Self { role: Role::Tool, content: result.to_string(), tool_call_id: Some(tool_call_id.to_string()), name: Some(name.to_string()), tool_calls: None }
}
}
// ─── 工具调用 ───
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
#[serde(rename = "type")]
pub call_type: String, // "function"
pub function: ToolFunctionCall,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolFunctionCall {
pub name: String,
pub arguments: String, // JSON string
}
// ─── 流式响应块 ───
#[derive(Debug, Clone)]
pub enum StreamChunk {
/// 文本片段 delta
Text(String),
/// 思考/推理内容(DeepSeek thinking
Reasoning(String),
/// 完成(携带完整文本,以及可选的工具调用)
Done {
text: String,
reasoning: Option<String>,
tool_calls: Option<Vec<ToolCall>>,
usage: Option<Usage>,
},
/// 错误
Error(String),
}
// ─── Token 用量 ───
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
#[serde(default)]
pub prompt_cache_hit_tokens: u32,
#[serde(default)]
pub prompt_cache_miss_tokens: u32,
}
// ─── 对话配置 ───
#[derive(Debug, Clone)]
pub struct ConversationConfig {
pub system_prompt: String,
pub model: String,
pub temperature: f32,
pub max_tokens: u32,
pub thinking: bool,
/// LLM function calling 的工具定义(JSON 数组)
pub tools: Option<Vec<serde_json::Value>>,
}
impl Default for ConversationConfig {
fn default() -> Self {
Self {
system_prompt: "You are a concise and helpful assistant replying in Chinese. \
Keep replies practical and natural."
.to_string(),
model: "deepseek-v4-flash".to_string(),
temperature: 0.7,
max_tokens: 4096,
thinking: true,
tools: None,
}
}
}
+635 -2
View File
@@ -1,3 +1,636 @@
fn main() {
println!("Hello, world!");
mod cli;
mod config;
mod context;
mod db;
mod llm;
mod scheduler;
mod state;
mod tools;
mod wechat;
use clap::Parser;
use cli::{Cli, Commands};
use context::MemoryStore;
use config::AppConfig;
use context::types::ChatSession as _;
use db::Database;
use llm::{Conversation, ConversationConfig, ToolExecutor, Usage, DEFAULT_SYSTEM_PROMPT};
use std::sync::Arc;
use tools::approval::ApprovalManager;
use tools::registry::{ExecutionContext, SkillRegistry};
use tracing::{error, info};
use tracing_subscriber::EnvFilter;
use wechat::client::WeChatClient;
#[tokio::main]
async fn main() {
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("info")),
)
.init();
dotenvy::dotenv().ok();
let config = AppConfig::from_file("config.json");
// 初始化数据库(如果配置了 DATABASE_URL
let database = match Database::connect().await {
Ok(db) => {
info!("✅ 数据库连接成功");
Some(Arc::new(db))
}
Err(e) => {
info!("数据库不可用,使用文件存储: {}", e);
None
}
};
// 长期记忆存储
let memory_store = Arc::new(MemoryStore::new(database.as_ref().map(|db| Arc::new(db.pool().clone()))));
memory_store.load().await;
// 文件状态管理器(降级方案)
let file_state = state::StateManager::new(&config.storage.state_dir);
let cli = Cli::parse();
match cli.command {
Commands::Login { timeout } => {
cmd_login(timeout, &database, &file_state).await;
}
Commands::Listen {
llm: enable_llm,
echo,
} => {
cmd_listen(enable_llm, echo, &database, &file_state, &config, &memory_store).await;
}
Commands::Send {
to,
text,
context_token,
} => {
cmd_send(&to, &text, context_token.as_deref(), &database, &file_state).await;
}
Commands::Whoami => {
cmd_whoami(&database, &file_state).await;
}
Commands::Usage { since, until, model } => {
cmd_usage(&database, since, until, model).await;
}
}
}
// ─── 命令实现 ───
async fn cmd_login(
timeout_secs: u64,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
) {
let client = WeChatClient::new(None);
match client
.login(|url| println!("二维码链接: {}", url), timeout_secs)
.await
{
Ok(result) => {
let auth = state::AuthState {
token: result.token.clone(),
account_id: result.account_id.clone(),
base_url: result.base_url.clone(),
};
// 优先存数据库
if let Some(db) = database {
if let Err(e) = db::models::save_auth(db.pool(), &auth).await {
error!("保存 auth 到数据库失败: {}", e);
file_state.save_auth(&auth);
} else {
info!("认证信息已保存到数据库");
// 清理文件状态(迁移完成)
file_state.clear_auth();
}
} else {
file_state.save_auth(&auth);
}
println!("\n✅ 登录成功!");
println!(" 账号: {}", result.account_id);
println!(" Token: {}...", &result.token[..std::cmp::min(16, result.token.len())]);
println!(" API: {}", result.base_url);
}
Err(e) => error!("登录失败: {}", e),
}
}
async fn cmd_listen(
enable_llm: bool,
echo: bool,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
config: &AppConfig,
memory_store: &Arc<MemoryStore>,
) {
// 从数据库或文件加载认证
let auth = 'load: {
if let Some(db) = database {
if let Some(a) = db::models::load_auth(db.pool()).await {
break 'load Some(a);
}
// 尝试从文件迁移
if let Some(fa) = file_state.load_auth() {
let a = state::AuthState {
token: fa.token,
account_id: fa.account_id,
base_url: fa.base_url,
};
let _ = db::models::save_auth(db.pool(), &a).await;
break 'load Some(a);
}
} else if let Some(a) = file_state.load_auth() {
break 'load Some(a);
}
break 'load None;
};
let auth = match auth {
Some(a) => a,
None => {
error!("未登录,请先执行 login 命令");
return;
}
};
let client = WeChatClient::new(Some(auth.base_url.clone()));
client
.set_auth(&auth.token, &auth.account_id, &auth.base_url)
.await;
// 从文件恢复 runtime 状态(get_updates_buf 暂时保留文件存储)
if let Some(r) = file_state.load_runtime() {
client.set_updates_buf(&r.get_updates_buf).await;
}
if let Err(e) = client.notify_start().await {
error!("注册监听器失败: {}", e);
return;
}
// 加载系统 Skills
let skill_registry = match SkillRegistry::load("resources/skills", "skills").await {
Ok(r) => {
info!("已加载 {} 个技能", r.count());
Some(Arc::new(r))
}
Err(e) => {
info!("技能加载结果: {:?}", e);
None
}
};
// 审批管理器
let approval_manager = database.as_ref().map(|db| {
Arc::new(ApprovalManager::new(Some(Arc::new(db.pool().clone()))))
});
let conversation = if enable_llm {
let sys_prompt = std::env::var("WEIXIN_LLM_SYSTEM_PROMPT")
.unwrap_or_else(|_| DEFAULT_SYSTEM_PROMPT.to_string());
let model = config.llm.deepseek.model.clone();
let mut cfg = ConversationConfig {
system_prompt: sys_prompt,
model,
..Default::default()
};
// 构建 LLM tools 列表(skills + 上下文工具)
let mut tools_list: Vec<serde_json::Value> = Vec::new();
if let Some(ref registry) = skill_registry {
for spec in registry.list_specs() {
tools_list.push(serde_json::json!({
"type": "function",
"function": {
"name": spec.name,
"description": spec.description,
"parameters": spec.parameters,
}
}));
}
}
// 上下文工具(长期记忆 + 摘要)
tools_list.push(serde_json::json!({
"type": "function",
"function": {
"name": "read_memories",
"description": "读取用户的长期记忆(跨会话持久保存的个人信息、偏好、约定等)",
"parameters": { "type": "object", "properties": {} }
}
}));
tools_list.push(serde_json::json!({
"type": "function",
"function": {
"name": "write_memory",
"description": "写入一条用户的长期记忆(当用户提到重要偏好、个人信息、约定时调用)",
"parameters": {
"type": "object",
"properties": {
"content": { "type": "string", "description": "记忆内容" }
},
"required": ["content"]
}
}
}));
tools_list.push(serde_json::json!({
"type": "function",
"function": {
"name": "read_summaries",
"description": "读取历史会话摘要(包括因上下文压缩或空闲超时产生的摘要)",
"parameters": { "type": "object", "properties": {} }
}
}));
cfg.tools = Some(tools_list);
// 创建 Conversation
let mut conv = match Conversation::new(cfg) {
Ok(c) => {
info!("LLM 已初始化: {} / {}", c.provider_name(), c.model());
c
}
Err(e) => {
error!("初始化 LLM 失败: {}", e);
return;
}
};
// 构建工具执行器
let session = conv.session();
let registry_opt = skill_registry.clone();
let approval = approval_manager.clone();
let memory_store = memory_store.clone();
let send_wechat: tools::registry::WechatSender = {
let client = client.clone();
Arc::new(move |user_id: &str, text: &str| {
let client = client.clone();
let uid = user_id.to_string();
let msg = text.to_string();
Box::pin(async move {
client.send_text(&uid, &msg, None).await.map(|_| ()).map_err(|e| e)
})
})
};
let executor: ToolExecutor = Arc::new(move |name: &str, args_json: &str| {
let reg = registry_opt.clone();
let approval = approval.clone();
let send = send_wechat.clone();
let session = session.clone();
let memory_store = memory_store.clone();
let n = name.to_string();
let args = args_json.to_string();
Box::pin(async move {
match n.as_str() {
"read_memories" => return Ok(memory_store.read().await),
"write_memory" => {
let params: serde_json::Value = serde_json::from_str(&args).unwrap_or_default();
let content = params.get("content").and_then(|v| v.as_str()).unwrap_or("");
if content.is_empty() { return Ok("请提供 content 参数".to_string()); }
return Ok(memory_store.write(content).await);
}
"read_summaries" => return Ok(context::tools::read_summaries(&session).await),
_ => {}
}
if let Some(ref reg) = reg {
let params: serde_json::Value = serde_json::from_str(&args).unwrap_or(serde_json::json!({}));
let mut ctx = ExecutionContext::new("wechat_user");
if let Some(a) = approval { ctx = ctx.with_approval(a); }
ctx = ctx.with_wechat(send);
let result = reg.execute(&n, params, &ctx).await;
Ok(result.output)
} else {
Ok(format!("未知工具: {}", n))
}
})
});
conv.set_tool_executor(executor);
Some(conv)
} else {
None
};
let account_id = auth.account_id.clone();
// 启动调度器(如果数据库可用)
if let Some(sched_db) = database.as_ref().map(|db| db.pool().clone()) {
let sched_client = client.clone();
tokio::spawn(async move {
let scheduler = scheduler::Scheduler::new(Some(Arc::new(sched_db)));
scheduler.run(move |to: &str, _name: &str, msg: &str| {
let c = sched_client.clone();
let uid = to.to_string();
let text = msg.to_string();
tokio::task::spawn(async move {
if let Err(e) = c.send_text(&uid, &text, None).await {
tracing::error!("发送调度结果失败: {}", e);
}
});
Ok(())
}).await;
});
info!("定时任务调度器已启动");
}
info!("开始监听消息 (llm={}, echo={})...", enable_llm, echo);
println!("已开始监听消息,按 Ctrl+C 退出");
let shutdown = async { tokio::signal::ctrl_c().await.expect("Ctrl+C 失败") };
tokio::select! {
_ = shutdown => {
info!("收到退出信号,注销监听器...");
let _ = client.notify_stop().await;
file_state.save_runtime(&client.get_updates_buf().await);
info!("已退出");
}
_ = listen_loop(&client, enable_llm, echo, &account_id, database, file_state, conversation, approval_manager) => {}
}
}
async fn listen_loop(
client: &WeChatClient,
enable_llm: bool,
echo: bool,
account_id: &str,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
conversation: Option<Conversation>,
approval_manager: Option<Arc<ApprovalManager>>,
) {
loop {
match client.receive_messages().await {
Ok(resp) => {
// 持久化 runtime buf
if !resp.get_updates_buf.is_empty() {
file_state.save_runtime(&resp.get_updates_buf);
}
for msg in &resp.msgs {
if msg.msg_type != Some(wechat::types::WeixinMessage::TYPE_USER) {
continue;
}
let from = msg.from_user_id.as_deref().unwrap_or("unknown");
let text = msg.text_content().unwrap_or("(非文本消息)");
let ctx_token = msg.context_token.as_deref();
let msg_id = msg.message_id.unwrap_or(0).to_string();
info!("收到消息 from={}: {}", from, text);
// === 入库:收到的消息 ===
if let Some(db) = database {
let _ = db::models::insert_chat_record(
db.pool(),
"inbound",
from,
account_id,
text,
"wechat",
ctx_token,
&msg_id,
)
.await;
}
// Echo 模式
if echo {
match client.send_text(from, text, ctx_token).await {
Ok(sent_id) => {
info!("回显成功 msg_id={}", sent_id);
// === 入库:发送的回显 ===
if let Some(db) = database {
let _ = db::models::insert_chat_record(
db.pool(),
"outbound",
from,
account_id,
text,
"echo",
ctx_token,
&sent_id,
)
.await;
}
}
Err(e) => error!("回显失败: {}", e),
}
}
// 审批回复检查:是否匹配待审批的确认码
let is_approval_reply = if let Some(ref mgr) = approval_manager {
mgr.handle_reply(from, text).await.is_some()
} else {
false
};
if is_approval_reply {
info!("消息匹配审批确认码,不传给 LLM(已由审批系统处理)");
continue;
}
// LLM 回复
if enable_llm {
if let Some(ref conv) = conversation {
let reply_result = if conv.spec().tools.is_some() {
generate_reply_with_tools(conv, text.to_string(), database).await
} else {
generate_reply(conv, text.to_string(), database).await
};
match reply_result {
Ok(reply) => {
match client.send_text(from, &reply, ctx_token).await {
Ok(sent_id) => {
info!("LLM 回复成功 msg_id={}", sent_id);
if let Some(db) = database {
let _ = db::models::insert_chat_record(
db.pool(), "outbound", from, account_id,
&reply, "llm", ctx_token, &sent_id,
).await;
}
}
Err(e) => error!("发送 LLM 回复失败: {}", e),
}
}
Err(e) => error!("生成 LLM 回复失败: {}", e),
}
}
}
}
}
Err(e) => {
if !e.contains("超时") && !e.contains("timeout") {
error!("接收消息失败: {}", e);
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
}
}
}
}
}
async fn generate_reply(conv: &Conversation, user_text: String, db: &Option<Arc<Database>>) -> Result<String, String> {
let handle = conv.chat(user_text).await?;
if let Some(u) = handle.get_usage() {
info!("📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
u.total_tokens, u.prompt_tokens, u.completion_tokens,
u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens);
store_usage(db, "", conv.model(), conv.provider_name(), u).await;
}
handle.consume().await
}
async fn generate_reply_with_tools(
conv: &Conversation,
user_text: String,
db: &Option<Arc<Database>>,
) -> Result<String, String> {
let (reply, _used_tools, usage) = conv.chat_with_tools(user_text).await?;
if let Some(u) = &usage {
info!("📊 Token: {} 总 / {} 提示 / {} 生成 | 缓存命中 {} | 缓存未命中 {}",
u.total_tokens, u.prompt_tokens, u.completion_tokens,
u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens);
store_usage(db, "", conv.model(), conv.provider_name(), u).await;
}
Ok(reply)
}
async fn store_usage(db: &Option<Arc<Database>>, user_id: &str, model: &str, provider: &str, u: &Usage) {
if let Some(db) = db {
let _ = db::models::insert_llm_usage(
db.pool(), user_id, model, provider,
u.prompt_tokens, u.completion_tokens,
u.prompt_cache_hit_tokens, u.prompt_cache_miss_tokens,
).await;
}
}
async fn cmd_whoami(
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
) {
let auth: Option<state::AuthState> = if let Some(db) = database {
db::models::load_auth(db.pool()).await
} else {
file_state.load_auth()
};
match auth {
Some(a) => {
println!("账号: {}", a.account_id);
println!("Token: {}...", &a.token[..std::cmp::min(16, a.token.len())]);
println!("API: {}", a.base_url);
if database.is_some() {
println!("存储: PostgreSQL");
} else {
println!("存储: 文件 ({})", file_state.state_dir().display());
}
}
None => println!("未登录,请先执行 login 命令"),
}
}
async fn cmd_usage(
database: &Option<Arc<Database>>,
since: Option<String>,
until: Option<String>,
model: Option<String>,
) {
let db = match database {
Some(d) => d,
None => { println!("数据库未配置,无法查询用量"); return; }
};
let since_dt = since.and_then(|s| {
chrono::DateTime::parse_from_rfc3339(&s).ok().map(|d| d.with_timezone(&chrono::Utc))
});
let until_dt = until.and_then(|s| {
chrono::DateTime::parse_from_rfc3339(&s).ok().map(|d| d.with_timezone(&chrono::Utc))
});
let model_ref = model.as_deref();
match db::models::query_llm_usage_stats(db.pool(), since_dt, until_dt, model_ref).await {
Ok(stats) => {
let hit_rate = if stats.total_cache_hit + stats.total_cache_miss > 0 {
stats.total_cache_hit as f64 / (stats.total_cache_hit + stats.total_cache_miss) as f64 * 100.0
} else { 0.0 };
println!("📊 LLM Token 使用统计");
println!("{:=<40}", "");
println!("调用次数: {}", stats.total_calls);
println!("Prompt Tokens: {} ({:.1}K)", stats.total_prompt_tokens, stats.total_prompt_tokens as f64 / 1000.0);
println!("生成 Tokens: {} ({:.1}K)", stats.total_completion_tokens, stats.total_completion_tokens as f64 / 1000.0);
println!("总 Tokens: {} ({:.1}K)", stats.total_tokens, stats.total_tokens as f64 / 1000.0);
println!("缓存命中: {} ({:.1}%)", stats.total_cache_hit, hit_rate);
println!("缓存未命中: {} ({:.1}%)", stats.total_cache_miss, 100.0 - hit_rate);
}
Err(e) => println!("查询失败: {}", e),
}
}
async fn cmd_send(
to: &str,
text: &str,
context_token: Option<&str>,
database: &Option<Arc<Database>>,
file_state: &state::StateManager,
) {
let auth: Option<state::AuthState> = if let Some(db) = database {
db::models::load_auth(db.pool()).await
} else {
file_state.load_auth()
};
let auth = match auth {
Some(a) => a,
None => {
error!("未登录,请先执行 login 命令");
return;
}
};
let base_url = auth.base_url.clone();
let account_id = auth.account_id.clone();
let client = WeChatClient::new(Some(auth.base_url));
client
.set_auth(&auth.token, &auth.account_id, &base_url)
.await;
match client.send_text(to, text, context_token).await {
Ok(msg_id) => {
println!("✅ 消息已发送, msg_id={}", msg_id);
// 入库:发送的消息
if let Some(db) = database {
let _ = db::models::insert_chat_record(
db.pool(),
"outbound",
to,
&account_id,
text,
"manual",
context_token,
&msg_id,
)
.await;
}
}
Err(e) => error!("发送失败: {}", e),
}
}
+167
View File
@@ -0,0 +1,167 @@
use chrono::Utc;
use sqlx::PgPool;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::interval;
/// 定时任务调度器 —— 轮询 PostgreSQL 中到期的 scheduled_tasks
pub struct Scheduler {
pool: Option<Arc<PgPool>>,
}
impl Scheduler {
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
Self { pool }
}
/// 启动调度循环(在 listen_loop 中作为 tokio::spawn 调用)
/// 每 5 秒检查一次是否有到期任务,执行后通过 callback 发送结果
pub async fn run(
&self,
mut on_task: impl FnMut(&str, &str, &str) -> Result<(), String>,
) {
let pool = match self.pool.as_ref() {
Some(p) => p.clone(),
None => {
tracing::info!("调度器:数据库未配置,跳过");
return;
}
};
let mut tick = interval(Duration::from_secs(5));
tracing::info!("调度器已启动(每 5s 检查)");
loop {
tick.tick().await;
let tasks = match fetch_due_tasks(&pool).await {
Ok(t) => t,
Err(e) => {
tracing::warn!("调度器查询失败: {}", e);
continue;
}
};
for task in tasks {
tracing::info!("⏰ 执行定时任务: {}", task.name);
// 执行任务
let result = execute_task(&task).await;
// 更新任务状态
let _ = update_task_status(&pool, &task.id, &result).await;
// 通知用户
let notify_msg = format!(
"⏰ 定时任务: {}\n结果: {}",
task.name,
result
);
if let Err(e) = on_task(&task.user_id, &task.name, &notify_msg) {
tracing::error!("发送定时任务通知失败: {}", e);
}
tracing::info!("✅ 定时任务完成: {} → {}", task.name, result);
}
}
}
}
#[derive(Debug)]
struct DueTask {
id: uuid::Uuid,
name: String,
user_id: String,
command: String,
shell: String,
cwd: String,
}
async fn fetch_due_tasks(pool: &PgPool) -> Result<Vec<DueTask>, String> {
let rows = sqlx::query_as::<_, (uuid::Uuid, String, String, String, String, String)>(
r#"
SELECT id, name, user_id, command, shell, cwd
FROM scheduled_tasks
WHERE enabled = true AND next_run_at <= NOW()
ORDER BY next_run_at ASC
LIMIT 10
"#,
)
.fetch_all(pool)
.await
.map_err(|e| format!("查询到期任务失败: {}", e))?;
Ok(rows
.into_iter()
.map(|(id, name, user_id, command, shell, cwd)| DueTask {
id,
name,
user_id,
command,
shell,
cwd,
})
.collect())
}
async fn execute_task(task: &DueTask) -> String {
let shell = if task.shell.is_empty() {
"/bin/bash"
} else {
&task.shell
};
let cwd = if task.cwd.is_empty() {
std::env::current_dir()
.map(|p| p.to_string_lossy().to_string())
.unwrap_or_default()
} else {
task.cwd.clone()
};
let result = tokio::process::Command::new(shell)
.arg("-c")
.arg(&task.command)
.current_dir(&cwd)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.output()
.await;
match result {
Ok(output) => {
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
if output.status.success() {
if stdout.is_empty() { "执行成功(无输出)".to_string() } else { stdout }
} else {
format!(
"退出码 {}: {}",
output.status.code().unwrap_or(-1),
if stderr.is_empty() { &stdout } else { &stderr }
)
}
}
Err(e) => format!("执行失败: {}", e),
}
}
async fn update_task_status(pool: &PgPool, task_id: &uuid::Uuid, result: &str) -> Result<(), String> {
// 更新 last_run_at 和 next_run_at
sqlx::query(
r#"
UPDATE scheduled_tasks
SET last_run_at = NOW(),
next_run_at = NOW() + (interval_seconds * INTERVAL '1 second'),
last_status = $2
WHERE id = $1
"#,
)
.bind(task_id)
.bind(result.chars().take(500).collect::<String>())
.execute(pool)
.await
.map_err(|e| format!("更新任务状态失败: {}", e))?;
Ok(())
}
+96
View File
@@ -0,0 +1,96 @@
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
/// 认证状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthState {
pub token: String,
pub account_id: String,
pub base_url: String,
}
/// 运行时状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RuntimeState {
pub get_updates_buf: String,
}
/// 状态管理器
pub struct StateManager {
state_dir: PathBuf,
}
impl StateManager {
pub fn new(state_dir: impl Into<PathBuf>) -> Self {
Self {
state_dir: state_dir.into(),
}
}
fn auth_file(&self) -> PathBuf {
self.state_dir.join("auth.json")
}
fn runtime_file(&self) -> PathBuf {
self.state_dir.join("runtime.json")
}
fn ensure_dir(&self) {
std::fs::create_dir_all(&self.state_dir).ok();
}
pub fn state_dir(&self) -> &std::path::Path {
&self.state_dir
}
pub fn save_auth(&self, auth: &AuthState) {
self.ensure_dir();
let json = serde_json::to_string_pretty(auth).expect("序列化 auth 失败");
if let Err(e) = std::fs::write(self.auth_file(), &json) {
tracing::error!("保存 auth 状态失败: {}", e);
}
// 权限
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(self.auth_file(), std::fs::Permissions::from_mode(0o600)).ok();
}
}
pub fn load_auth(&self) -> Option<AuthState> {
let path = self.auth_file();
if !path.exists() {
return None;
}
let json = std::fs::read_to_string(&path).ok()?;
serde_json::from_str(&json).ok()
}
pub fn save_runtime(&self, buf: &str) {
self.ensure_dir();
let state = RuntimeState {
get_updates_buf: buf.to_string(),
};
let json = serde_json::to_string_pretty(&state).expect("序列化 runtime 失败");
if let Err(e) = std::fs::write(self.runtime_file(), &json) {
tracing::error!("保存 runtime 状态失败: {}", e);
}
}
pub fn load_runtime(&self) -> Option<RuntimeState> {
let path = self.runtime_file();
if !path.exists() {
return None;
}
let json = std::fs::read_to_string(&path).ok()?;
serde_json::from_str(&json).ok()
}
/// 删除 auth 文件(迁移到数据库后清理)
pub fn clear_auth(&self) {
let path = self.auth_file();
if path.exists() {
std::fs::remove_file(&path).ok();
}
}
}
+328
View File
@@ -0,0 +1,328 @@
# iAs 工具系统设计
## 核心理念
**一切皆 Skill。**
没有"内置工具"和"外部技能"的区别——所有工具都是 SKILL.md,区别仅在来源:
- **系统 skills** — `resources/skills/` 目录,随 iAs 分发,只读
- **用户 skills** — `skills/` 项目目录,用户自己添加
加载时自动合并,同名时用户 skills 覆盖系统 skills。
---
## 架构
```
┌──────────────────────────────────────────────────────────┐
│ LLM │
│ generate_reply() ←→ tool_choice → tool_execute(result) │
└─────────────────────┬────────────────────────────────────┘
┌──────────────────────────────────────────────────────────┐
│ SkillRegistry │
│ │
│ ┌──────────────────┐ ┌────────────────────────────┐ │
│ │ 系统 Skills │ │ 用户 Skills │ │
│ │ resources/skills/ │ │ skills/ │ │
│ │ │ │ │ │
│ │ • datetime │ │ • query_qweather/ │ │
│ │ • weather │ │ • search_tavily/ │ │
│ │ • search_web │ │ • start_pc/ │ │
│ │ • email │ │ • ... │ │
│ │ • shell_exec │ │ │ │
│ │ • schedule │ │ │ │
│ └────────┬──────────┘ └────────────┬──────────────┘ │
│ │ │ │
│ ▼ ▼ │
│ ┌──────────────────────────────────────────────────┐ │
│ │ Approval Gate │ │
│ │ risk == High → 确认码审批 / Low → 自动执行 │ │
│ │ 审批过程对 LLM 完全透明 │ │
│ └──────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────────┘
```
---
## 核心类型
```rust
/// 风险等级(仅两级)
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RiskLevel {
Low, // 只读/安全,自动执行
High, // 写操作/有风险,需微信确认码审批
}
/// 技能来源
#[derive(Debug, Clone, PartialEq)]
pub enum SkillSource {
System, // resources/skills/
User, // skills/
}
/// 技能元数据(从 SKILL.md 解析)
#[derive(Debug, Clone, Deserialize)]
pub struct SkillSpec {
pub name: String,
pub description: String,
pub risk_level: RiskLevel,
pub parameters: serde_json::Value, // JSON Schema
pub timeout_secs: u64,
}
/// 技能(加载到内存后)
pub struct Skill {
pub spec: SkillSpec,
pub source: SkillSource,
pub dir: PathBuf,
pub execute_command: String,
}
```
---
## SKILL.md 格式
每个技能是一个目录,包含 `SKILL.md` 文件:
```markdown
# Get Current Datetime
获取当前日期时间(北京时间 UTC+8)。
## Risk Level
Low
## Parameters
\`\`\`json
{
"type": "object",
"properties": {
"format": {
"type": "string",
"description": "时间格式: full | date | time",
"enum": ["full", "date", "time"],
"default": "full"
}
}
}
\`\`\`
## Execute
\`\`\`bash
#!/usr/bin/env bash
scripts/datetime.sh
\`\`\`
```
---
## 内置 Skills 一览
随 iAs 分发的系统 skills
| Skill | 风险 | 说明 |
|-------|------|------|
| `get_current_datetime` | Low | 获取当前时间(UTC+8 |
| `query_weather` | Low | 和风天气查询 |
| `search_web` | Low | Tavily 网络搜索 |
| `email` | High | 收发邮件 |
| `execute_shell` | High | 执行 shell 命令 |
| `list_scheduled_tasks` | Low | 列出定时任务 |
| `manage_scheduled_task` | High | 增删改定时任务 |
---
## SkillRegistry
```rust
pub struct SkillRegistry {
skills: HashMap<String, Arc<Skill>>,
}
impl SkillRegistry {
pub async fn load(system_dir: &str, user_dir: &str) -> Result<Self, LoadError>;
pub fn list_specs(&self) -> Vec<&SkillSpec>;
pub fn list_specs_by_risk(&self, level: RiskLevel) -> Vec<&SkillSpec>;
pub fn get(&self, name: &str) -> Option<&Skill>;
/// 执行技能(含审批流程)
pub async fn execute(
&self,
name: &str,
params: serde_json::Value,
ctx: &ExecutionContext,
) -> Result<SkillResult, SkillError>;
}
pub struct ExecutionContext {
pub user_id: String,
pub db: Option<Arc<Database>>,
/// 发送微信消息的回调
pub send_wechat: Option<Box<dyn Fn(&str, &str) -> Result<(), String> + Send + Sync>>,
/// 等待用户回复的回调(返回用户消息文本)
pub wait_for_reply: Option<Box<dyn Fn(&str, Duration) -> Result<Option<String>, String> + Send + Sync>>,
}
```
---
## 审批流程(透明模式)
**核心原则:审批过程对 LLM 完全不可见。** LLM 只看到最终结果。
```
LLM 调用 High 风险技能
SkillRegistry::execute()
├── 查找 Skill
├── risk_level == High
└── YES ──────────────────────────────────────────┐
│ │
▼ │
┌──────────────────────────────────────────┐ │
│ 等待阶段(LLM 不感知) │ │
│ │ │
│ ① 生成 6 位确认码(如 482731) │ │
│ ② SHA-256 哈希存入 pending_approvals │ │
│ ③ 通过微信发送确认消息给用户: │ │
│ │ │
│ ┌─────────────────────────────┐ │ │
│ │ ⚠️ 操作确认 │ │ │
│ │ │ │ │
│ │ 技能:send_email │ │ │
│ │ 参数:收件人 xxx@xxx.com │ │ │
│ │ │ │ │
│ │ 确认码:482731 │ │ │
│ │ 回复确认码以继续 │ │ │
│ │ 回复 0 或「取消」以取消 │ │ │
│ │ (5分钟内有效,最多3次尝试) │ │ │
│ └─────────────────────────────┘ │ │
│ │ │
│ ④ 等待用户回复 │ │
│ │ │ │
│ ├── 回复正确确认码 ──┐ │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ ┌──────────────┐ │ │
│ │ │ ⑤ 执行技能 │ │ │
│ │ │ ⑥ LLM 看到 │ │ │
│ │ │ 执行结果 │ │ │
│ │ └──────────────┘ │ │
│ │ │ │
│ ├── 回复错误确认码 │ │
│ │ → "确认码错误,还剩 N 次" │ │
│ │ → 重试(最多 3 次) │ │
│ │ → 3 次都错 │ │
│ │ → LLM 看到: │ │
│ │ "用户拒绝了你的调用: │ │
│ │ send_email" │ │
│ │ │ │
│ ├── 回复 0 / 「取消」 │ │
│ │ → LLM 看到: │ │
│ │ "用户拒绝了你的调用: │ │
│ │ send_email" │ │
│ │ │ │
│ └── 5 分钟超时 │ │
│ → LLM 看到: │ │
│ "用户没有确认操作: │ │
│ send_email" │ │
└──────────────────────────────────────────┘
```
### LLM 视角(两种可能)
```
✅ 通过:
LLM: 调用 send_email({to: "xxx", subject: "..."})
系统: { success: true, message: "邮件已发送" }
❌ 拒绝:
LLM: 调用 send_email({to: "xxx", subject: "..."})
系统: { success: false, error: "用户拒绝了你的调用:send_email" }
```
LLM 不知道确认码的存在,不知道审批过程。它只知道自己调用了工具,然后得到了结果或拒绝。
---
## 数据库表
```sql
-- 待审批的工具调用
CREATE TABLE pending_approvals (
id UUID PRIMARY KEY,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
expires_at TIMESTAMPTZ NOT NULL, -- 5分钟后过期
user_id TEXT NOT NULL,
skill_name TEXT NOT NULL,
params JSONB NOT NULL,
code_hash TEXT NOT NULL, -- 确认码 SHA-256(不存明文)
attempts_left INTEGER NOT NULL DEFAULT 3, -- 剩余尝试次数
status TEXT NOT NULL DEFAULT 'pending', -- pending | approved | rejected | expired
result JSONB,
consumed_at TIMESTAMPTZ
);
```
---
## Skills 目录结构
```
resources/skills/ # 系统内置 Skills(随 iAs 分发)
├── get_current_datetime/
│ ├── SKILL.md
│ └── scripts/datetime.sh
├── query_weather/
│ ├── SKILL.md
│ └── scripts/weather.sh
├── search_web/
│ ├── SKILL.md
│ └── scripts/search.sh
├── email/
│ ├── SKILL.md
│ └── scripts/email.sh
├── execute_shell/
│ ├── SKILL.md
│ └── scripts/shell.sh
├── list_scheduled_tasks/
│ ├── SKILL.md
│ └── scripts/list.sh
└── manage_scheduled_task/
├── SKILL.md
└── scripts/manage.sh
skills/ # 用户 Skills(用户自行添加)
├── query_qweather/
│ ├── SKILL.md
│ └── scripts/query.sh
├── search_tavily/
│ ├── SKILL.md
│ └── scripts/search.sh
└── start_pc/
├── SKILL.md
└── scripts/wake.sh
```
---
## 实施步骤
| 步骤 | 内容 | 预估 |
|------|------|------|
| Step 1 | SkillSpec/Skill/SkillRegistry 核心类型 + 注册表 | 半天 |
| Step 2 | SKILL.md 解析器 + 目录扫描(system + user 合并) | 半天 |
| Step 3 | 技能执行器:子进程管理 + 参数 JSON 传递 + 结果解析 | 半天 |
| Step 4 | 系统 Skills 脚本实现(7 个 bash 脚本) | 1 天 |
| Step 5 | 审批系统:pending_approvals 表 + 确认码生成/验证 + 微信透明通知 | 1 天 |
| Step 6 | LLM function calling 集成 + 审计日志 | 1 天 |
+154
View File
@@ -0,0 +1,154 @@
use rand::Rng;
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, oneshot};
/// 审批决策
#[derive(Debug, Clone, PartialEq)]
pub enum ApprovalDecision {
Approved,
Rejected,
Expired,
}
/// 内存中待审批的条目
struct PendingEntry {
code_hash: String,
skill_name: String,
attempts_left: i32,
/// 通知等待方
sender: oneshot::Sender<ApprovalDecision>,
}
/// 审批管理器
pub struct ApprovalManager {
pool: Option<Arc<PgPool>>,
/// user_id → 该用户的待审批队列
pending: Arc<Mutex<HashMap<String, Vec<PendingEntry>>>>,
}
impl ApprovalManager {
pub fn new(pool: Option<Arc<PgPool>>) -> Self {
Self {
pool,
pending: Arc::new(Mutex::new(HashMap::new())),
}
}
/// 创建审批请求
/// 返回 (确认码明文, 接收器) — 调用方 await 接收器等待用户确认
pub async fn create(
&self,
user_id: &str,
skill_name: &str,
) -> Result<(String, oneshot::Receiver<ApprovalDecision>), String> {
let code: u32 = rand::thread_rng().gen_range(100000..999999);
let code_str = code.to_string();
let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes()));
let (tx, rx) = oneshot::channel();
// 入库
if let Some(ref pool) = self.pool {
let id = uuid::Uuid::new_v4();
let expires_at = chrono::Utc::now() + chrono::Duration::minutes(5);
let params = serde_json::json!({"name": skill_name});
let _ = sqlx::query(
r#"
INSERT INTO pending_approvals (id, expires_at, user_id, skill_name, params, code_hash, status)
VALUES ($1, $2, $3, $4, $5, $6, 'pending')
"#,
)
.bind(id)
.bind(expires_at)
.bind(user_id)
.bind(skill_name)
.bind(&params)
.bind(&code_hash)
.execute(pool.as_ref())
.await;
}
// 注册到内存
let mut map = self.pending.lock().await;
map.entry(user_id.to_string())
.or_default()
.push(PendingEntry {
code_hash,
skill_name: skill_name.to_string(),
attempts_left: 3,
sender: tx,
});
Ok((code_str, rx))
}
/// 处理用户回复(由消息循环调用)
/// 匹配成功时返回 Some(技能名),匹配失败返回 None
pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<String> {
let mut map = self.pending.lock().await;
let entries = map.get_mut(user_id)?;
if entries.is_empty() {
return None;
}
let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
let reply_trimmed = reply.trim();
// 检查取消
if reply_trimmed == "0" || reply_trimmed == "取消" {
let entry = entries.remove(0);
let name = entry.skill_name.clone();
let _ = entry.sender.send(ApprovalDecision::Rejected);
return Some(name);
}
// 检查确认码
let pos = entries.iter().position(|e| e.code_hash == reply_hash);
if let Some(idx) = pos {
let entry = entries.remove(idx);
let name = entry.skill_name.clone();
let _ = entry.sender.send(ApprovalDecision::Approved);
return Some(name);
}
// 确认码错误
if let Some(entry) = entries.first_mut() {
entry.attempts_left -= 1;
if entry.attempts_left <= 0 {
let entry = entries.remove(0);
let name = entry.skill_name.clone();
let _ = entry.sender.send(ApprovalDecision::Rejected);
return Some(name);
}
}
None
}
/// 清理过期(由定时任务调用)
pub async fn clean_expired(&self) {
let mut map = self.pending.lock().await;
for entries in map.values_mut() {
entries.retain(|e| !e.sender.is_closed());
}
map.retain(|_, v| !v.is_empty());
if let Some(ref pool) = self.pool {
let _ = sqlx::query(
"UPDATE pending_approvals SET status = 'expired' WHERE status = 'pending' AND expires_at < NOW()",
)
.execute(pool.as_ref())
.await;
}
}
/// 获取某个用户的待审批条数
pub async fn pending_count(&self, user_id: &str) -> usize {
let map = self.pending.lock().await;
map.get(user_id).map(|v| v.len()).unwrap_or(0)
}
}
+4
View File
@@ -0,0 +1,4 @@
pub mod approval;
pub mod parser;
pub mod registry;
pub mod skill;
+279
View File
@@ -0,0 +1,279 @@
use crate::tools::skill::{Skill, SkillSource, SkillSpec};
use std::fs;
use std::path::Path;
/// 解析 SKILL.md 文件,返回 Skill
///
/// SKILL.md 格式:
/// ```markdown
/// # Skill Name
/// Description text...
///
/// ## Risk Level
/// Low | High
///
/// ## Parameters
/// ```json
/// { ... }
/// ```
///
/// ## Execute
/// ```bash
/// command to run
/// ```
/// ```
pub fn parse_skill(skill_dir: &Path) -> Result<Skill, String> {
let md_path = skill_dir.join("SKILL.md");
let content =
fs::read_to_string(&md_path).map_err(|e| format!("读取 {} 失败: {}", md_path.display(), e))?;
let lines: Vec<&str> = content.lines().collect();
let mut name = String::new();
let mut description = String::new();
let mut risk_level = String::new();
let mut params_raw = String::new();
let mut execute_raw = String::new();
let mut in_params = false;
let mut in_execute = false;
let mut params_fence = 0;
let mut execute_fence = 0;
let mut desc_lines: Vec<&str> = Vec::new();
let mut in_desc = false;
for (i, line) in lines.iter().enumerate() {
let trimmed = line.trim();
// 标题行 → 技能名
if trimmed.starts_with("# ") && !trimmed.starts_with("## ") && name.is_empty() {
name = trimmed.trim_start_matches("# ").trim().to_string();
continue;
}
// 二级标题
if trimmed.starts_with("## ") {
let section = trimmed.trim_start_matches("## ").trim().to_lowercase();
match section.as_str() {
"risk level" | "risk_level" | "risk" => {
in_desc = false;
in_params = false;
in_execute = false;
}
"parameters" | "param" => {
in_desc = false;
in_params = true;
in_execute = false;
}
"execute" | "execution" | "command" | "run" => {
in_desc = false;
in_params = false;
in_execute = true;
params_fence = 0;
execute_fence = 0;
}
_ => {
in_desc = false;
in_params = false;
in_execute = false;
}
}
continue;
}
// 描述:标题和 Risk Level 之间的文本
if !in_params && !in_execute && !trimmed.starts_with("#") && !trimmed.is_empty() {
if name.is_empty() {
continue;
}
// 检查是否已进入 Risk Level 段
let is_risk_line = trimmed.to_lowercase().starts_with("low")
|| trimmed.to_lowercase().starts_with("high");
if risk_level.is_empty() && is_risk_line {
risk_level = trimmed.to_string();
continue;
}
if risk_level.is_empty() {
desc_lines.push(trimmed);
}
continue;
}
// Risk Level 值
if !in_params && !in_execute && risk_level.is_empty() {
let lower = trimmed.to_lowercase();
if lower == "low" || lower == "high" {
risk_level = trimmed.to_string();
continue;
}
}
// Parameters: JSON 代码块
if in_params {
if trimmed.starts_with("```") {
params_fence += 1;
if params_fence == 2 {
in_params = false;
}
continue;
}
if params_fence == 1 {
if !params_raw.is_empty() {
params_raw.push('\n');
}
params_raw.push_str(trimmed);
}
continue;
}
// Execute: shell 命令
if in_execute {
if trimmed.starts_with("```") {
execute_fence += 1;
if execute_fence == 2 {
in_execute = false;
}
continue;
}
if execute_fence == 1 {
// 跳过 shebang 行
if trimmed.starts_with("#!/") {
continue;
}
if !execute_raw.is_empty() {
execute_raw.push('\n');
}
execute_raw.push_str(trimmed);
}
continue;
}
}
// 验证
if name.is_empty() {
return Err(format!("{}: 未找到技能名称 (# Title)", md_path.display()));
}
description = desc_lines.join(" ").trim().to_string();
let risk = match risk_level.to_lowercase().trim() {
"high" => crate::tools::skill::RiskLevel::High,
_ => crate::tools::skill::RiskLevel::Low,
};
// 解析参数 JSON
let parameters: serde_json::Value = if params_raw.is_empty() {
serde_json::json!({
"type": "object",
"properties": {}
})
} else {
serde_json::from_str(&params_raw).unwrap_or_else(|_| {
serde_json::json!({
"type": "object",
"properties": {}
})
})
};
// 执行命令(去掉首尾空白)
let execute_command = execute_raw.trim().to_string();
if execute_command.is_empty() {
return Err(format!("{}: 未找到 Execute 命令", md_path.display()));
}
let spec = SkillSpec {
name,
description,
risk_level: risk,
parameters,
timeout_secs: 30,
};
Ok(Skill::new(
spec,
// 先设成 User,由调用方修正
SkillSource::User,
skill_dir.to_path_buf(),
execute_command,
))
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#[test]
fn test_load_all_system_skills() {
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(crate::tools::registry::SkillRegistry::load(
"resources/skills",
"/nonexistent",
));
match result {
Ok(registry) => {
assert_eq!(registry.count(), 8, "应该有 8 个系统技能");
let names: Vec<&str> = registry.list_specs().iter().map(|s| s.name.as_str()).collect();
println!("已加载技能: {:?}", names);
assert!(names.contains(&"manage_memos"));
assert!(names.contains(&"get_current_datetime"));
assert!(names.contains(&"query_weather"));
assert!(names.contains(&"search_web"));
assert!(names.contains(&"email"));
assert!(names.contains(&"execute_shell"));
assert!(names.contains(&"list_scheduled_tasks"));
assert!(names.contains(&"manage_scheduled_task"));
}
Err(errors) => {
for e in &errors {
eprintln!("加载失败: {}", e);
}
panic!("技能加载失败");
}
}
}
#[test]
fn test_parse_basic_skill() {
let dir = std::env::temp_dir().join("ias-test-skill");
fs::create_dir_all(&dir).ok();
let mut f = fs::File::create(dir.join("SKILL.md")).unwrap();
write!(
f,
r#"# Get Current Datetime
获取当前日期时间(北京时间 UTC+8)。
## Risk Level
Low
## Parameters
```json
{{
"type": "object",
"properties": {{
"format": {{
"type": "string",
"description": "full | date | time"
}}
}}
}}
```
## Execute
```bash
date '+%Y-%m-%d %H:%M:%S'
```
"#
)
.unwrap();
let skill = parse_skill(&dir).unwrap();
assert_eq!(skill.spec.name, "Get Current Datetime");
assert!(skill.spec.description.contains("日期时间"));
assert_eq!(skill.spec.risk_level, crate::tools::skill::RiskLevel::Low);
assert!(!skill.execute_command.is_empty());
assert!(skill.execute_command.contains("date"));
fs::remove_dir_all(&dir).ok();
}
}
+311
View File
@@ -0,0 +1,311 @@
use super::approval::{ApprovalDecision, ApprovalManager};
use crate::tools::parser::parse_skill;
use crate::tools::skill::*;
use std::future::Future;
use std::pin::Pin;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use tokio::process::Command;
/// 技能注册表
pub struct SkillRegistry {
skills: HashMap<String, Arc<Skill>>,
}
impl SkillRegistry {
pub fn new() -> Self {
Self {
skills: HashMap::new(),
}
}
/// 从 system_dir 和 user_dir 加载所有技能
/// system_dir: resources/skills/(随 iAs 分发)
/// user_dir: skills/(用户自行添加)
/// 同名时 user 覆盖 system
pub async fn load(system_dir: &str, user_dir: &str) -> Result<Self, Vec<String>> {
let mut registry = Self::new();
let mut errors = Vec::new();
// 加载系统 skills
if Path::new(system_dir).exists() {
registry.load_from_dir(system_dir, SkillSource::System, &mut errors);
}
// 加载用户 skills
if Path::new(user_dir).exists() {
registry.load_from_dir(user_dir, SkillSource::User, &mut errors);
}
if errors.is_empty() {
Ok(registry)
} else {
Err(errors)
}
}
fn load_from_dir(&mut self, dir: &str, source: SkillSource, errors: &mut Vec<String>) {
let dir_path = Path::new(dir);
let entries = match std::fs::read_dir(dir_path) {
Ok(e) => e,
Err(e) => {
errors.push(format!("读取目录 {} 失败: {}", dir, e));
return;
}
};
for entry in entries.flatten() {
let path = entry.path();
if !path.is_dir() {
continue;
}
let skill_md = path.join("SKILL.md");
if !skill_md.exists() {
continue;
}
match parse_skill(&path) {
Ok(mut skill) => {
skill.source = source.clone();
let name = skill.spec.name.clone();
// user 同名覆盖 system
if source == SkillSource::User && self.skills.contains_key(&name) {
tracing::info!("用户技能覆盖系统技能: {}", name);
}
self.skills.insert(name, Arc::new(skill));
}
Err(e) => {
errors.push(format!("{}: {}", path.display(), e));
}
}
}
}
/// 列出所有技能 specs(给 LLM function calling
pub fn list_specs(&self) -> Vec<&SkillSpec> {
let mut specs: Vec<&SkillSpec> = self.skills.values().map(|s| &s.spec).collect();
specs.sort_by(|a, b| a.name.cmp(&b.name));
specs
}
/// 按风险等级过滤
pub fn list_specs_by_risk(&self, level: RiskLevel) -> Vec<&SkillSpec> {
self.skills
.values()
.filter(|s| s.spec.risk_level == level)
.map(|s| &s.spec)
.collect()
}
/// 按名称查找技能
pub fn get(&self, name: &str) -> Option<&Skill> {
self.skills.get(name).map(|s| s.as_ref())
}
/// 获取技能数量
pub fn count(&self) -> usize {
self.skills.len()
}
// ─── 执行 ───
/// 执行技能(含审批流程)
pub async fn execute(
&self,
name: &str,
params: serde_json::Value,
ctx: &ExecutionContext,
) -> SkillResult {
let skill = match self.skills.get(name) {
Some(s) => s.clone(),
None => {
return SkillResult::error(format!("技能不存在: {}", name));
}
};
// High 风险 → 审批流程(透明模式)
if skill.spec.risk_level.is_high() {
match approval_flow(&skill, ctx).await {
ApprovalResult::Approved => {
// 继续执行
}
ApprovalResult::Rejected => {
return SkillResult::rejected(name);
}
ApprovalResult::Expired => {
return SkillResult::expired(name);
}
ApprovalResult::Error(msg) => {
return SkillResult::error(msg);
}
}
}
// 执行技能
run_skill(&skill, params, ctx).await
}
}
// ─── 审批流程 ───
enum ApprovalResult {
Approved,
Rejected,
Expired,
Error(String),
}
async fn approval_flow(skill: &Skill, ctx: &ExecutionContext) -> ApprovalResult {
let manager = match ctx.approval_manager.as_ref() {
Some(m) => m.clone(),
None => return ApprovalResult::Approved, // 测试模式
};
// 创建审批请求(生成确认码 + 注册通道)
let (code, rx) = match manager.create(&ctx.user_id, &skill.spec.name).await {
Ok(pair) => pair,
Err(e) => return ApprovalResult::Error(e),
};
tracing::debug!("审批确认码已生成: {}...", &code[..2]);
// 发送确认消息给用户
let msg = format!(
"⚠️ 操作确认\n\n技能:{}\n\n确认码:{}\n\n回复确认码以继续\n回复 0 或「取消」以取消\n(5分钟内有效,最多3次尝试)",
skill.spec.name, code
);
if let Some(ref send) = ctx.send_wechat {
if let Err(e) = send(&ctx.user_id, &msg).await {
return ApprovalResult::Error(format!("发送确认消息失败: {}", e));
}
}
// 等待用户确认(通过 oneshot channel,由消息循环驱动)
match tokio::time::timeout(Duration::from_secs(300), rx).await {
Ok(Ok(ApprovalDecision::Approved)) => ApprovalResult::Approved,
Ok(Ok(ApprovalDecision::Rejected)) => ApprovalResult::Rejected,
Ok(Ok(ApprovalDecision::Expired)) => ApprovalResult::Expired,
Ok(Err(_)) => ApprovalResult::Error("审批通道异常".to_string()),
Err(_) => ApprovalResult::Expired,
}
}
// ─── 技能执行器 ───
async fn run_skill(skill: &Skill, params: serde_json::Value, _ctx: &ExecutionContext) -> SkillResult {
let timeout = Duration::from_secs(skill.spec.timeout_secs);
let cmd_str = &skill.execute_command;
let parts: Vec<&str> = cmd_str.split_whitespace().collect();
if parts.is_empty() {
return SkillResult::error("执行命令为空");
}
let mut cmd = Command::new(parts[0]);
for arg in &parts[1..] {
cmd.arg(arg);
}
cmd.current_dir(&skill.dir);
tracing::debug!(
"执行命令: {:?} (cwd={:?})",
&cmd_str, &skill.dir
);
// 通过 stdin 传入参数
let params_json = serde_json::to_string(&params).unwrap_or_default();
cmd.stdin(std::process::Stdio::piped());
cmd.stdout(std::process::Stdio::piped());
cmd.stderr(std::process::Stdio::piped());
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => return SkillResult::error(format!("启动进程失败: {}", e)),
};
// 写入参数到 stdin
if let Some(mut stdin) = child.stdin.take() {
use tokio::io::AsyncWriteExt;
let _ = stdin.write_all(params_json.as_bytes()).await;
// 关闭 stdin 让进程知道输入结束
drop(stdin);
}
// 等待完成(带超时)
let result = tokio::time::timeout(timeout, child.wait_with_output()).await;
match result {
Ok(Ok(output)) => {
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
let exit_code = output.status.code().unwrap_or(-1);
tracing::debug!(
"🖥️ {} exit={} out={:.200} err={:.100}",
skill.spec.name, exit_code,
if stdout.is_empty() { "(空)" } else { &stdout },
if stderr.is_empty() { "(空)" } else { &stderr }
);
if output.status.success() {
// 尝试解析为 JSON
if let Ok(json) = serde_json::from_str::<serde_json::Value>(&stdout) {
SkillResult::ok_with_data(stdout, json)
} else {
SkillResult::ok(stdout)
}
} else {
let msg = if !stderr.is_empty() {
format!("退出码 {}: {}", output.status.code().unwrap_or(-1), stderr)
} else {
format!("退出码 {}", output.status.code().unwrap_or(-1))
};
SkillResult::error(msg)
}
}
Ok(Err(e)) => SkillResult::error(format!("进程错误: {}", e)),
Err(_) => SkillResult::error(format!("执行超时 ({}s)", skill.spec.timeout_secs)),
}
}
// ─── 执行上下文 ───
/// 发送微信消息的回调: (user_id, text) -> Result(异步)
pub type WechatSender =
Arc<dyn Fn(&str, &str) -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>;
pub struct ExecutionContext {
pub user_id: String,
pub db: Option<Arc<sqlx::PgPool>>,
/// 审批管理器
pub approval_manager: Option<Arc<ApprovalManager>>,
/// 发送微信消息的回调
pub send_wechat: Option<WechatSender>,
}
impl ExecutionContext {
pub fn new(user_id: impl Into<String>) -> Self {
Self {
user_id: user_id.into(),
db: None,
approval_manager: None,
send_wechat: None,
}
}
pub fn with_db(mut self, db: Arc<sqlx::PgPool>) -> Self {
self.db = Some(db);
self
}
pub fn with_approval(mut self, mgr: Arc<ApprovalManager>) -> Self {
self.approval_manager = Some(mgr);
self
}
pub fn with_wechat(mut self, send: WechatSender) -> Self {
self.send_wechat = Some(send);
self
}
}
+156
View File
@@ -0,0 +1,156 @@
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
// ─── 风险等级 ───
#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RiskLevel {
Low, // 只读/安全,自动执行
High, // 写操作/有风险,需微信确认码审批
}
impl RiskLevel {
pub fn is_high(&self) -> bool {
matches!(self, RiskLevel::High)
}
}
// ─── 技能来源 ───
#[derive(Debug, Clone, PartialEq)]
pub enum SkillSource {
System, // resources/skills/
User, // skills/
}
impl SkillSource {
pub fn label(&self) -> &str {
match self {
SkillSource::System => "系统",
SkillSource::User => "用户",
}
}
}
// ─── 技能元数据(从 SKILL.md 解析) ───
#[derive(Debug, Clone, Deserialize)]
pub struct SkillSpec {
pub name: String,
pub description: String,
#[serde(default = "default_risk_level")]
pub risk_level: RiskLevel,
#[serde(default)]
pub parameters: serde_json::Value,
#[serde(default = "default_timeout")]
pub timeout_secs: u64,
}
fn default_risk_level() -> RiskLevel {
RiskLevel::Low
}
fn default_timeout() -> u64 {
30
}
// ─── 技能(加载到内存后) ───
#[derive(Debug, Clone)]
pub struct Skill {
pub spec: SkillSpec,
pub source: SkillSource,
pub dir: PathBuf,
pub execute_command: String,
}
impl Skill {
pub fn new(spec: SkillSpec, source: SkillSource, dir: PathBuf, execute_command: String) -> Self {
Self {
spec,
source,
dir,
execute_command,
}
}
}
// ─── 执行上下文 ───
pub type WechatSender = Arc<dyn Fn(&str, &str) -> Result<(), String> + Send + Sync>;
pub type ReplyWaiter = Arc<dyn Fn(&str, Duration) -> Result<Option<String>, String> + Send + Sync>;
/// LLM 能看到的结果
#[derive(Debug, Clone, Serialize)]
pub struct SkillResult {
pub success: bool,
pub output: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub data: Option<serde_json::Value>,
}
impl SkillResult {
pub fn ok(output: impl Into<String>) -> Self {
Self {
success: true,
output: output.into(),
data: None,
}
}
pub fn ok_with_data(output: impl Into<String>, data: serde_json::Value) -> Self {
Self {
success: true,
output: output.into(),
data: Some(data),
}
}
pub fn rejected(skill_name: &str) -> Self {
Self {
success: false,
output: format!("用户拒绝了你的调用:{}", skill_name),
data: None,
}
}
pub fn expired(skill_name: &str) -> Self {
Self {
success: false,
output: format!("用户没有确认操作:{}", skill_name),
data: None,
}
}
pub fn error(msg: impl Into<String>) -> Self {
Self {
success: false,
output: msg.into(),
data: None,
}
}
}
// ─── 错误 ───
#[derive(Debug)]
pub enum SkillError {
NotFound(String),
Execution(String),
Timeout(String),
InvalidParams(String),
}
impl std::fmt::Display for SkillError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SkillError::NotFound(name) => write!(f, "技能不存在: {}", name),
SkillError::Execution(msg) => write!(f, "执行失败: {}", msg),
SkillError::Timeout(name) => write!(f, "执行超时: {}", name),
SkillError::InvalidParams(msg) => write!(f, "参数无效: {}", msg),
}
}
}
+436
View File
@@ -0,0 +1,436 @@
use crate::wechat::types::*;
use base64::Engine;
use reqwest::Client as HttpClient;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tracing::{debug, info, warn};
/// 微信 iLink Bot API 客户端
#[derive(Clone)]
pub struct WeChatClient {
http: HttpClient,
base_url: String,
token: Arc<Mutex<String>>,
account_id: Arc<Mutex<String>>,
/// get_updates_buf,用于长轮询的上下文缓冲
updates_buf: Arc<Mutex<String>>,
}
impl WeChatClient {
/// iLink App ID(从 package.json 的 ilink_appid 字段,硬编码或环境变量)
const ILINK_APP_ID: &'static str = "";
/// 默认 bot_type
const DEFAULT_BOT_TYPE: &'static str = "3";
/// 默认 API 地址
const DEFAULT_BASE_URL: &'static str = "https://ilinkai.weixin.qq.com";
/// 二维码固定 API 地址
const FIXED_BASE_URL: &'static str = "https://ilinkai.weixin.qq.com";
/// 版本号编码:0x00MMNNPP
const CLIENT_VERSION: u32 = 0x000100_00; // 1.0.0
pub fn new(base_url: Option<String>) -> Self {
let base = base_url.unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string());
Self {
http: HttpClient::builder()
.timeout(Duration::from_secs(60))
.build()
.expect("创建 HTTP client 失败"),
base_url: base,
token: Arc::new(Mutex::new(String::new())),
account_id: Arc::new(Mutex::new(String::new())),
updates_buf: Arc::new(Mutex::new(String::new())),
}
}
// ─── 公共方法 ───
/// 获取登录二维码链接
pub async fn get_qrcode(&self) -> Result<(String, String), String> {
let url = format!(
"{}/ilink/bot/get_bot_qrcode?bot_type={}",
Self::FIXED_BASE_URL,
Self::DEFAULT_BOT_TYPE
);
let resp: QRCodeResponse = self
.http
.get(&url)
.headers(Self::common_headers())
.send()
.await
.map_err(|e| format!("请求二维码失败: {}", e))?
.json()
.await
.map_err(|e| format!("解析二维码响应失败: {}", e))?;
Ok((resp.qrcode, resp.qrcode_img_content))
}
/// 轮询扫码状态
pub async fn poll_qr_status(
&self,
qrcode: &str,
base_url: &str,
) -> Result<StatusResponse, String> {
let url = format!(
"{}/ilink/bot/get_qrcode_status?qrcode={}",
base_url, qrcode
);
let resp_result = self
.http
.get(&url)
.headers(Self::common_headers())
.timeout(Duration::from_secs(35))
.send()
.await;
match resp_result {
Ok(response) => {
let text = response.text().await
.map_err(|e| format!("读取响应体失败: {}", e))?;
serde_json::from_str::<StatusResponse>(&text)
.map_err(|e| format!("解析状态响应失败: {}", e))
}
Err(e) if e.is_timeout() => {
// 超时是正常的,返回 wait 让调用者继续轮询
Ok(StatusResponse {
status: "wait".to_string(),
bot_token: None,
ilink_bot_id: None,
baseurl: None,
ilink_user_id: None,
redirect_host: None,
})
}
Err(e) => {
warn!("轮询二维码状态网络错误: {},继续等待", e);
Ok(StatusResponse {
status: "wait".to_string(),
bot_token: None,
ilink_bot_id: None,
baseurl: None,
ilink_user_id: None,
redirect_host: None,
})
}
}
}
/// 执行完整扫码登录流程
pub async fn login(
&self,
on_qrcode: impl Fn(&str),
timeout_secs: u64,
) -> Result<LoginResult, String> {
info!("开始微信扫码登录...");
let (qrcode, qrcode_img_content) = self.get_qrcode().await?;
// 显示二维码链接
on_qrcode(&qrcode_img_content);
// 生成二维码终端显示
if let Ok(code) = qrcode::QrCode::new(qrcode_img_content.as_bytes()) {
let svg = code.render::<qrcode::render::unicode::Dense1x2>()
.dark_color(qrcode::render::unicode::Dense1x2::Dark)
.light_color(qrcode::render::unicode::Dense1x2::Light)
.build();
println!("\n{}", svg);
}
println!("请使用微信扫描上方二维码登录");
println!("二维码链接(如二维码无法显示):{}", qrcode_img_content);
println!();
let deadline = tokio::time::Instant::now() + Duration::from_secs(timeout_secs);
let mut scanned = false;
let mut current_base = Self::FIXED_BASE_URL.to_string();
while tokio::time::Instant::now() < deadline {
let status = self.poll_qr_status(&qrcode, &current_base).await?;
match status.status.as_str() {
"wait" => {
print!(".");
std::io::Write::flush(&mut std::io::stdout()).ok();
}
"scaned" => {
if !scanned {
println!("\n👀 已扫码,请在微信上确认登录...");
scanned = true;
}
}
"scaned_but_redirect" => {
if let Some(host) = &status.redirect_host {
current_base = format!("https://{}", host);
info!("IDC 重定向到: {}", current_base);
}
}
"confirmed" => {
let token = status
.bot_token
.ok_or_else(|| "登录成功但未返回 bot_token".to_string())?;
let account_id = status
.ilink_bot_id
.ok_or_else(|| "登录成功但未返回 ilink_bot_id".to_string())?;
let user_id = status.ilink_user_id.unwrap_or_default();
let base_url = status.baseurl.unwrap_or_else(|| Self::DEFAULT_BASE_URL.to_string());
// 保存到内存状态
*self.token.lock().await = token.clone();
*self.account_id.lock().await = account_id.clone();
info!("✅ 登录成功! bot_id={}", account_id);
return Ok(LoginResult {
token,
account_id,
base_url,
user_id,
});
}
"expired" => {
return Err("二维码已过期".to_string());
}
_ => {
warn!("未知状态: {}", status.status);
}
}
tokio::time::sleep(Duration::from_secs(1)).await;
}
Err("登录超时".to_string())
}
/// 注册监听器(通知微信服务端本客户端已启动)
pub async fn notify_start(&self) -> Result<(), String> {
let body = serde_json::json!({ "base_info": BaseInfo::new() });
let url = format!("{}/ilink/bot/msg/notifystart", self.base_url);
let token = self.token.lock().await.clone();
let _resp: NotifyResp = self
.post_json(&url, &body, Some(&token))
.await?
.json()
.await
.map_err(|e| format!("解析 notifyStart 响应失败: {}", e))?;
info!("监听器已注册");
Ok(())
}
/// 注销监听器
pub async fn notify_stop(&self) -> Result<(), String> {
let body = serde_json::json!({ "base_info": BaseInfo::new() });
let url = format!("{}/ilink/bot/msg/notifystop", self.base_url);
let token = self.token.lock().await.clone();
let _resp: NotifyResp = self
.post_json(&url, &body, Some(&token))
.await?
.json()
.await
.map_err(|e| format!("解析 notifyStop 响应失败: {}", e))?;
info!("监听器已注销");
Ok(())
}
/// 长轮询接收消息
pub async fn receive_messages(&self) -> Result<GetUpdatesResp, String> {
let buf = self.updates_buf.lock().await.clone();
let body = GetUpdatesReq {
get_updates_buf: buf.clone(),
base_info: BaseInfo::new(),
};
let url = format!("{}/ilink/bot/getupdates", self.base_url);
let token = self.token.lock().await.clone();
let resp: GetUpdatesResp = match self
.post_json(&url, &body, Some(&token))
.await
{
Ok(resp) => resp.json().await.map_err(|e| format!("解析 getUpdates 响应失败: {}", e))?,
Err(e) => {
// 如果超时,返回空响应
if e.contains("超时") || e.contains("timeout") {
return Ok(GetUpdatesResp {
ret: 0,
errcode: None,
errmsg: None,
msgs: vec![],
get_updates_buf: buf.clone(),
longpolling_timeout_ms: None,
});
}
return Err(e);
}
};
// 更新 buf
if !resp.get_updates_buf.is_empty() {
*self.updates_buf.lock().await = resp.get_updates_buf.clone();
}
Ok(resp)
}
/// 发送文本消息
pub async fn send_text(
&self,
to: &str,
text: &str,
context_token: Option<&str>,
) -> Result<String, String> {
let client_id = uuid::Uuid::new_v4().to_string();
let msg = WeixinMessage {
from_user_id: Some(String::new()),
to_user_id: Some(to.to_string()),
client_id: Some(client_id.clone()),
msg_type: Some(WeixinMessage::TYPE_BOT),
message_state: Some(2), // FINISH
item_list: Some(vec![MessageItem::text(text)]),
context_token: context_token.map(|s| s.to_string()),
..Default::default()
};
let body = SendMessageReq { msg };
let url = format!("{}/ilink/bot/sendmessage", self.base_url);
let token = self.token.lock().await.clone();
let _resp_text = self
.post_json(&url, &body, Some(&token))
.await?
.text()
.await
.map_err(|e| format!("发送消息网络错误: {}", e))?;
Ok(client_id)
}
/// 设置 auth 状态(从持久化恢复)
pub async fn set_auth(&self, token: &str, account_id: &str, base_url: &str) {
*self.token.lock().await = token.to_string();
*self.account_id.lock().await = account_id.to_string();
if !base_url.is_empty() {
// 只更新 base_url 如果传入了非空值
}
}
/// 获取当前 token(用于持久化)
pub async fn get_token(&self) -> String {
self.token.lock().await.clone()
}
/// 获取当前 account_id(用于持久化)
pub async fn get_account_id(&self) -> String {
self.account_id.lock().await.clone()
}
/// 获取当前 updates_buf(用于持久化)
pub async fn get_updates_buf(&self) -> String {
self.updates_buf.lock().await.clone()
}
/// 设置 get_updates_buf(从持久化恢复)
pub async fn set_updates_buf(&self, buf: &str) {
*self.updates_buf.lock().await = buf.to_string();
}
// ─── 内部方法 ───
fn common_headers() -> reqwest::header::HeaderMap {
let mut headers = reqwest::header::HeaderMap::new();
headers.insert(
"Content-Type",
reqwest::header::HeaderValue::from_static("application/json"),
);
if !Self::ILINK_APP_ID.is_empty() {
headers.insert(
"iLink-App-Id",
reqwest::header::HeaderValue::from_static(Self::ILINK_APP_ID),
);
}
headers.insert(
"iLink-App-ClientVersion",
reqwest::header::HeaderValue::from_str(&Self::CLIENT_VERSION.to_string())
.unwrap(),
);
// X-WECHAT-UIN: random uint32 -> decimal -> base64
let uin = rand::random::<u32>();
let uin_b64 = base64::engine::general_purpose::STANDARD
.encode(uin.to_string());
headers.insert(
"X-WECHAT-UIN",
reqwest::header::HeaderValue::from_str(&uin_b64).unwrap(),
);
headers
}
fn auth_headers(token: &str) -> reqwest::header::HeaderMap {
let mut headers = Self::common_headers();
if !token.is_empty() {
headers.insert(
reqwest::header::AUTHORIZATION,
reqwest::header::HeaderValue::from_str(&format!("Bearer {}", token))
.unwrap(),
);
headers.insert(
"AuthorizationType",
reqwest::header::HeaderValue::from_static("ilink_bot_token"),
);
}
headers
}
async fn post_json(
&self,
url: &str,
body: &impl serde::Serialize,
token: Option<&str>,
) -> Result<reqwest::Response, String> {
let body_str = serde_json::to_string(body).map_err(|e| format!("序列化请求体失败: {}", e))?;
debug!("POST {} body_len={}", url, body_str.len());
let mut headers = if let Some(t) = token {
Self::auth_headers(t)
} else {
Self::common_headers()
};
// 设置 Content-Length
headers.insert(
"Content-Length",
reqwest::header::HeaderValue::from_str(&body_str.len().to_string())
.unwrap(),
);
self.http
.post(url)
.headers(headers)
.body(body_str)
.send()
.await
.map_err(|e| {
if e.is_timeout() {
format!("请求超时: {}", url)
} else {
format!("请求失败 ({}): {}", url, e)
}
})
}
}
+2
View File
@@ -0,0 +1,2 @@
pub mod client;
pub mod types;
+333
View File
@@ -0,0 +1,333 @@
use serde::{Deserialize, Serialize};
// ─── 基础类型 ───
/// 公共请求元数据
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct BaseInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub channel_version: Option<String>,
}
impl BaseInfo {
pub fn new() -> Self {
Self {
channel_version: Some(env!("CARGO_PKG_VERSION").to_string()),
}
}
}
// ─── 消息类型 ───
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TextItem {
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CDNMedia {
#[serde(skip_serializing_if = "Option::is_none")]
pub encrypt_query_param: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aes_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub encrypt_type: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub full_url: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageItem {
#[serde(skip_serializing_if = "Option::is_none")]
pub media: Option<CDNMedia>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_media: Option<CDNMedia>,
#[serde(skip_serializing_if = "Option::is_none")]
pub aeskey: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub mid_size: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_size: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_height: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_width: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub hd_size: Option<i64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VoiceItem {
#[serde(skip_serializing_if = "Option::is_none")]
pub media: Option<CDNMedia>,
#[serde(skip_serializing_if = "Option::is_none")]
pub encode_type: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub bits_per_sample: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sample_rate: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub playtime: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileItem {
#[serde(skip_serializing_if = "Option::is_none")]
pub media: Option<CDNMedia>,
#[serde(skip_serializing_if = "Option::is_none")]
pub file_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub md5: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub len: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VideoItem {
#[serde(skip_serializing_if = "Option::is_none")]
pub media: Option<CDNMedia>,
#[serde(skip_serializing_if = "Option::is_none")]
pub video_size: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub play_length: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub video_md5: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_media: Option<CDNMedia>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_size: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_height: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thumb_width: Option<i32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RefMessage {
#[serde(skip_serializing_if = "Option::is_none")]
pub message_item: Option<Box<MessageItem>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
}
/// 消息条目(支持 text/image/voice/file/video
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageItem {
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
pub item_type: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub create_time_ms: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub update_time_ms: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub is_completed: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub msg_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ref_msg: Option<Box<RefMessage>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub text_item: Option<TextItem>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image_item: Option<ImageItem>,
#[serde(skip_serializing_if = "Option::is_none")]
pub voice_item: Option<VoiceItem>,
#[serde(skip_serializing_if = "Option::is_none")]
pub file_item: Option<FileItem>,
#[serde(skip_serializing_if = "Option::is_none")]
pub video_item: Option<VideoItem>,
}
impl MessageItem {
pub const TYPE_TEXT: i32 = 1;
pub const TYPE_IMAGE: i32 = 2;
pub const TYPE_VOICE: i32 = 3;
pub const TYPE_FILE: i32 = 4;
pub const TYPE_VIDEO: i32 = 5;
pub fn text(text: &str) -> Self {
Self {
item_type: Some(Self::TYPE_TEXT),
text_item: Some(TextItem {
text: Some(text.to_string()),
}),
..Default::default()
}
}
}
impl Default for MessageItem {
fn default() -> Self {
Self {
item_type: None,
create_time_ms: None,
update_time_ms: None,
is_completed: None,
msg_id: None,
ref_msg: None,
text_item: None,
image_item: None,
voice_item: None,
file_item: None,
video_item: None,
}
}
}
// ─── 微信消息 ───
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WeixinMessage {
#[serde(skip_serializing_if = "Option::is_none")]
pub seq: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub message_id: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub from_user_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub to_user_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub client_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub create_time_ms: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub update_time_ms: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub delete_time_ms: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
#[serde(rename = "message_type", skip_serializing_if = "Option::is_none")]
pub msg_type: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub message_state: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub item_list: Option<Vec<MessageItem>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub context_token: Option<String>,
}
impl Default for WeixinMessage {
fn default() -> Self {
Self {
seq: None,
message_id: None,
from_user_id: None,
to_user_id: None,
client_id: None,
create_time_ms: None,
update_time_ms: None,
delete_time_ms: None,
session_id: None,
group_id: None,
msg_type: None,
message_state: None,
item_list: None,
context_token: None,
}
}
}
impl WeixinMessage {
pub const TYPE_NONE: i32 = 0;
pub const TYPE_USER: i32 = 1;
pub const TYPE_BOT: i32 = 2;
/// 获取文本内容(如果有纯文本 item)
pub fn text_content(&self) -> Option<&str> {
self.item_list
.as_ref()?
.iter()
.find(|item| item.item_type == Some(MessageItem::TYPE_TEXT))
.and_then(|item| item.text_item.as_ref())
.and_then(|t| t.text.as_deref())
}
}
// ─── API 请求/响应 ───
/// 获取二维码响应
#[derive(Debug, Deserialize)]
pub struct QRCodeResponse {
pub qrcode: String,
pub qrcode_img_content: String,
}
/// 扫码状态响应
#[derive(Debug, Deserialize)]
pub struct StatusResponse {
pub status: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub bot_token: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ilink_bot_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub baseurl: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ilink_user_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub redirect_host: Option<String>,
}
/// GetUpdates 请求
#[derive(Debug, Serialize)]
pub struct GetUpdatesReq {
pub get_updates_buf: String,
pub base_info: BaseInfo,
}
/// GetUpdates 响应
#[derive(Debug, Deserialize)]
pub struct GetUpdatesResp {
#[serde(default)]
pub ret: i32,
#[serde(skip_serializing_if = "Option::is_none")]
pub errcode: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub errmsg: Option<String>,
#[serde(default)]
pub msgs: Vec<WeixinMessage>,
#[serde(default)]
pub get_updates_buf: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub longpolling_timeout_ms: Option<i64>,
}
/// SendMessage 请求
#[derive(Debug, Serialize)]
pub struct SendMessageReq {
pub msg: WeixinMessage,
}
/// NotifyStart/NotifyStop 响应
#[derive(Debug, Deserialize)]
pub struct NotifyResp {
#[serde(default)]
pub ret: i32,
#[serde(skip_serializing_if = "Option::is_none")]
pub errmsg: Option<String>,
}
/// 登录结果
#[derive(Debug, Clone)]
pub struct LoginResult {
pub token: String,
pub account_id: String,
pub base_url: String,
pub user_id: String,
}
/// 收到的消息事件
#[derive(Debug, Clone)]
pub struct MessageEvent {
pub from_user_id: String,
pub text: String,
pub context_token: Option<String>,
pub message_id: i64,
}