opensource/proxycast
3
0
Fork 0
mirror of https://github.com/aiclientproxy/proxycast.git synced 2026-10-28 03:43:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

release: v0.91.0

Browse Source
This commit is contained in:
coso
2026-03-19 19:05:53 +08:00
parent ad8427be93
commit a414f44872
175 changed files with 10467 additions and 11013 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/crates/services/src/aster_session_store.rs
+636 -165
View File
@@ -18,10 +18,12 @@ use aster::session::{
};
use async_trait::async_trait;
use chrono::Utc;
use lime_core::app_paths;
use lime_core::database::DbConnection;
use lime_core::workspace::WorkspaceManager;
use serde::de::DeserializeOwned;
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::path::{Path, PathBuf};
/// Lime 的 SessionStore 实现
///
@@ -36,6 +38,107 @@ impl LimeSessionStore {
Self { db }
}
pub fn load_extension_data_from_conn(
conn: &rusqlite::Connection,
session_id: &str,
) -> Result<ExtensionData> {
let extension_data_json: String = conn
.query_row(
"SELECT extension_data_json FROM agent_sessions WHERE id = ?1",
rusqlite::params![session_id],
|row| row.get(0),
)
.map_err(|e| anyhow!("读取 extension_data 失败: {e}"))?;
Ok(serde_json::from_str(&extension_data_json).unwrap_or_default())
}
pub fn load_extension_data_sync(db: &DbConnection, session_id: &str) -> Result<ExtensionData> {
let conn = db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
Self::load_extension_data_from_conn(&conn, session_id)
}
fn normalize_optional_text(value: Option<String>) -> Option<String> {
let value = value?;
let trimmed = value.trim();
if trimmed.is_empty() {
None
} else {
Some(trimmed.to_string())
}
}
fn parse_optional_json<T: DeserializeOwned>(raw: Option<String>) -> Option<T> {
raw.and_then(|text| serde_json::from_str(&text).ok())
}
fn resolve_session_type(raw: Option<String>, model: &str) -> SessionType {
let parsed_model = model.parse::<SessionType>().ok();
match raw
.as_deref()
.and_then(|value| value.parse::<SessionType>().ok())
{
Some(SessionType::User) if matches!(parsed_model, Some(parsed) if parsed != SessionType::User) => {
parsed_model.unwrap_or(SessionType::User)
}
Some(session_type) => session_type,
None => parsed_model.unwrap_or(SessionType::User),
}
}
fn default_model_name() -> String {
"agent:default".to_string()
}
fn insert_session_row(
conn: &rusqlite::Connection,
id: &str,
title: &str,
working_dir: &Path,
session_type: SessionType,
) -> Result<()> {
let now = Utc::now().to_rfc3339();
conn.execute(
"INSERT INTO agent_sessions (
id, model, system_prompt, title, created_at, updated_at, working_dir,
execution_strategy, session_type, user_set_name, extension_data_json
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
rusqlite::params![
id,
Self::default_model_name(),
None::<String>,
title,
now,
now,
working_dir.to_string_lossy().to_string(),
"react",
session_type.to_string(),
false,
serde_json::to_string(&ExtensionData::default())
.map_err(|e| anyhow!("序列化 extension_data 失败: {e}"))?,
],
)
.map_err(|e| anyhow!("创建会话失败: {e}"))?;
Ok(())
}
fn ensure_session_row(conn: &rusqlite::Connection, session_id: &str) -> Result<()> {
let session_exists: bool = conn
.query_row(
"SELECT 1 FROM agent_sessions WHERE id = ?",
[session_id],
|_| Ok(true),
)
.unwrap_or(false);
if session_exists {
return Ok(());
}
let working_dir = Self::resolve_session_working_dir(conn);
Self::insert_session_row(conn, session_id, "新对话", &working_dir, SessionType::User)
}
/// 将 Message 的 role 转换为字符串
/// 通过检查 Message::user() 和 Message::assistant() 的 role 来判断
fn message_role_to_string(message: &Message) -> String {
@@ -50,38 +153,23 @@ impl LimeSessionStore {
/// 解析会话 working_dir（优先默认 workspace，其次应用默认项目目录）
fn resolve_session_working_dir(conn: &rusqlite::Connection) -> PathBuf {
// 1) 优先使用默认 workspace（is_default = 1）
let default_workspace_path: Option<String> = conn
.query_row(
"SELECT root_path FROM workspaces WHERE is_default = 1 LIMIT 1",
[],
|row| row.get(0),
)
.ok();
if let Some(path) = default_workspace_path {
if !path.trim().is_empty() {
let pb = PathBuf::from(path);
return if pb.is_absolute() {
pb
} else {
std::env::current_dir()
.unwrap_or_else(|_| PathBuf::from("."))
.join(pb)
};
if let Some(path) = WorkspaceManager::get_default_root_path_from_conn(conn)
.ok()
.flatten()
{
let normalized = Self::normalize_working_dir(path);
if !normalized.as_os_str().is_empty() {
return normalized;
}
}
// 2) 回退到 ~/.lime/projects/default
if let Some(home) = dirs::home_dir() {
let fallback = home.join(".lime").join("projects").join("default");
if !fallback.exists() {
let _ = fs::create_dir_all(&fallback);
}
return fallback;
if let Ok(default_project_dir) = app_paths::resolve_default_project_dir() {
return default_project_dir;
}
// 3) 最终回退到进程当前目录
tracing::warn!(
"[SessionStore] 解析默认 working_dir 失败，已回退当前目录；建议检查 app_paths 配置"
);
std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
}
@@ -120,26 +208,8 @@ impl SessionStore for LimeSessionStore {
) -> Result<Session> {
let id = uuid::Uuid::new_v4().to_string();
let now = Utc::now();
let now_str = now.to_rfc3339();
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let type_str = session_type.to_string();
conn.execute(
"INSERT INTO agent_sessions (id, model, system_prompt, title, created_at, updated_at, working_dir)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
rusqlite::params![
id,
type_str,
None::<String>,
name,
now_str,
now_str,
working_dir.to_string_lossy().to_string()
],
)
.map_err(|e| anyhow!("创建会话失败: {e}"))?;
Self::insert_session_row(&conn, &id, &name, &working_dir, session_type)?;
Ok(Session {
id,
@@ -174,40 +244,17 @@ impl SessionStore for LimeSessionStore {
);
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
// 检查 session 是否存在
let session_exists: bool = conn
.query_row("SELECT 1 FROM agent_sessions WHERE id = ?", [id], |_| {
Ok(true)
})
.unwrap_or(false);
tracing::info!("[SessionStore] session_exists={}", session_exists);
// 如果不存在，自动创建
if !session_exists {
let now = Utc::now().to_rfc3339();
let working_dir = Self::resolve_session_working_dir(&conn);
conn.execute(
"INSERT INTO agent_sessions (id, model, system_prompt, title, created_at, updated_at, working_dir)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
rusqlite::params![
id,
"agent:default",
None::<String>,
"新对话",
now,
now,
working_dir.to_string_lossy().to_string()
],
)
.map_err(|e| anyhow!("自动创建会话失败: {e}"))?;
tracing::info!("[SessionStore] get_session 自动创建会话: {}", id);
}
Self::ensure_session_row(&conn, id)?;
tracing::info!("[SessionStore] get_session 已确保会话存在: {}", id);
let mut stmt = conn
.prepare(
"SELECT id, model, system_prompt, title, created_at, updated_at, working_dir
"SELECT id, model, system_prompt, title, created_at, updated_at, working_dir,
session_type, user_set_name, extension_data_json,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json, user_recipe_values_json,
provider_name, model_config_json
FROM agent_sessions WHERE id = ?",
)
.map_err(|e| anyhow!("准备查询失败: {e}"))?;
@@ -222,12 +269,47 @@ impl SessionStore for LimeSessionStore {
row.get::<_, String>(4)?,
row.get::<_, String>(5)?,
row.get::<_, Option<String>>(6)?,
row.get::<_, Option<String>>(7)?,
row.get::<_, bool>(8)?,
row.get::<_, String>(9)?,
row.get::<_, Option<i32>>(10)?,
row.get::<_, Option<i32>>(11)?,
row.get::<_, Option<i32>>(12)?,
row.get::<_, Option<i32>>(13)?,
row.get::<_, Option<i32>>(14)?,
row.get::<_, Option<i32>>(15)?,
row.get::<_, Option<String>>(16)?,
row.get::<_, Option<String>>(17)?,
row.get::<_, Option<String>>(18)?,
row.get::<_, Option<String>>(19)?,
row.get::<_, Option<String>>(20)?,
))
})
.map_err(|e| anyhow!("会话不存在: {e}"))?;
let (id, model, _system_prompt, title, created_at, updated_at, db_working_dir) =
session_row;
let (
id,
model,
_system_prompt,
title,
created_at,
updated_at,
db_working_dir,
session_type_raw,
user_set_name,
extension_data_json,
total_tokens,
input_tokens,
output_tokens,
accumulated_total_tokens,
accumulated_input_tokens,
accumulated_output_tokens,
schedule_id,
recipe_json,
user_recipe_values_json,
provider_name,
model_config_json,
) = session_row;
let created_at = chrono::DateTime::parse_from_rfc3339(&created_at)
.map(|dt| dt.with_timezone(&Utc))
@@ -236,7 +318,7 @@ impl SessionStore for LimeSessionStore {
.map(|dt| dt.with_timezone(&Utc))
.unwrap_or_else(|_| Utc::now());
let session_type = model.parse().unwrap_or(SessionType::User);
let session_type = Self::resolve_session_type(session_type_raw, &model);
let working_dir = Self::parse_session_working_dir(&conn, db_working_dir);
let conversation = if include_messages {
@@ -251,27 +333,29 @@ impl SessionStore for LimeSessionStore {
id: id.to_string(),
working_dir,
name: title.unwrap_or_else(|| "未命名会话".to_string()),
user_set_name: false,
user_set_name,
session_type,
created_at,
updated_at,
extension_data: ExtensionData::default(),
total_tokens: None,
input_tokens: None,
output_tokens: None,
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
schedule_id: None,
recipe: None,
user_recipe_values: None,
extension_data: serde_json::from_str(&extension_data_json).unwrap_or_default(),
total_tokens,
input_tokens,
output_tokens,
accumulated_total_tokens,
accumulated_input_tokens,
accumulated_output_tokens,
schedule_id,
recipe: Self::parse_optional_json(recipe_json),
user_recipe_values: Self::parse_optional_json(user_recipe_values_json),
conversation,
message_count,
provider_name: None,
model_config: match model.trim() {
"" | "agent:default" => None,
normalized => ModelConfig::new(normalized).ok(),
},
provider_name,
model_config: Self::parse_optional_json(model_config_json).or_else(|| {
match model.trim() {
"" | "agent:default" => None,
normalized => ModelConfig::new(normalized).ok(),
}
}),
})
}
@@ -282,35 +366,7 @@ impl SessionStore for LimeSessionStore {
);
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
// 检查会话是否存在，如果不存在则自动创建
let session_exists: bool = conn
.query_row(
"SELECT 1 FROM agent_sessions WHERE id = ?",
[session_id],
|_| Ok(true),
)
.unwrap_or(false);
if !session_exists {
let now = Utc::now().to_rfc3339();
let working_dir = Self::resolve_session_working_dir(&conn);
conn.execute(
"INSERT INTO agent_sessions (id, model, system_prompt, title, created_at, updated_at, working_dir)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
rusqlite::params![
session_id,
"agent:default",
None::<String>,
"新对话",
now,
now,
working_dir.to_string_lossy().to_string()
],
)
.map_err(|e| anyhow!("自动创建会话失败: {e}"))?;
tracing::info!("[SessionStore] 自动创建会话: {}", session_id);
}
Self::ensure_session_row(&conn, session_id)?;
let role = Self::message_role_to_string(message);
let content_json = serde_json::to_string(&message.content)
@@ -425,7 +481,12 @@ impl SessionStore for LimeSessionStore {
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let mut stmt = conn.prepare(
"SELECT id, model, system_prompt, title, created_at, updated_at, working_dir
"SELECT id, model, system_prompt, title, created_at, updated_at, working_dir,
session_type, user_set_name, extension_data_json,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
schedule_id, recipe_json, user_recipe_values_json,
provider_name, model_config_json
FROM agent_sessions ORDER BY updated_at DESC",
)?;
@@ -437,43 +498,106 @@ impl SessionStore for LimeSessionStore {
let created_at: String = row.get(4)?;
let updated_at: String = row.get(5)?;
let working_dir: Option<String> = row.get(6)?;
let session_type: Option<String> = row.get(7)?;
let user_set_name: bool = row.get(8)?;
let extension_data_json: String = row.get(9)?;
let total_tokens: Option<i32> = row.get(10)?;
let input_tokens: Option<i32> = row.get(11)?;
let output_tokens: Option<i32> = row.get(12)?;
let accumulated_total_tokens: Option<i32> = row.get(13)?;
let accumulated_input_tokens: Option<i32> = row.get(14)?;
let accumulated_output_tokens: Option<i32> = row.get(15)?;
let schedule_id: Option<String> = row.get(16)?;
let recipe_json: Option<String> = row.get(17)?;
let user_recipe_values_json: Option<String> = row.get(18)?;
let provider_name: Option<String> = row.get(19)?;
let model_config_json: Option<String> = row.get(20)?;
Ok((id, model, title, created_at, updated_at, working_dir))
Ok((
id,
model,
title,
created_at,
updated_at,
working_dir,
session_type,
user_set_name,
extension_data_json,
total_tokens,
input_tokens,
output_tokens,
accumulated_total_tokens,
accumulated_input_tokens,
accumulated_output_tokens,
schedule_id,
recipe_json,
user_recipe_values_json,
provider_name,
model_config_json,
))
})?
.filter_map(|r| r.ok())
.map(
|(id, model, title, created_at, updated_at, db_working_dir)| {
|(
id,
model,
title,
created_at,
updated_at,
db_working_dir,
session_type_raw,
user_set_name,
extension_data_json,
total_tokens,
input_tokens,
output_tokens,
accumulated_total_tokens,
accumulated_input_tokens,
accumulated_output_tokens,
schedule_id,
recipe_json,
user_recipe_values_json,
provider_name,
model_config_json,
)| {
let created_at = chrono::DateTime::parse_from_rfc3339(&created_at)
.map(|dt| dt.with_timezone(&Utc))
.unwrap_or_else(|_| Utc::now());
let updated_at = chrono::DateTime::parse_from_rfc3339(&updated_at)
.map(|dt| dt.with_timezone(&Utc))
.unwrap_or_else(|_| Utc::now());
let session_type = model.parse().unwrap_or(SessionType::User);
let session_type = Self::resolve_session_type(session_type_raw, &model);
let working_dir = Self::parse_session_working_dir(&conn, db_working_dir);
let message_count = self.count_messages(&conn, &id).unwrap_or(0);
Session {
id,
working_dir,
name: title.unwrap_or_else(|| "未命名会话".to_string()),
user_set_name: false,
user_set_name,
session_type,
created_at,
updated_at,
extension_data: ExtensionData::default(),
total_tokens: None,
input_tokens: None,
output_tokens: None,
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
schedule_id: None,
recipe: None,
user_recipe_values: None,
extension_data: serde_json::from_str(&extension_data_json)
.unwrap_or_default(),
total_tokens,
input_tokens,
output_tokens,
accumulated_total_tokens,
accumulated_input_tokens,
accumulated_output_tokens,
schedule_id,
recipe: Self::parse_optional_json(recipe_json),
user_recipe_values: Self::parse_optional_json(user_recipe_values_json),
conversation: None,
message_count: 0,
provider_name: None,
model_config: None,
message_count,
provider_name,
model_config: Self::parse_optional_json(model_config_json).or_else(|| {
match model.trim() {
"" | "agent:default" => None,
normalized => ModelConfig::new(normalized).ok(),
}
}),
}
},
)
@@ -504,10 +628,16 @@ impl SessionStore for LimeSessionStore {
let total_sessions: i64 =
conn.query_row("SELECT COUNT(*) FROM agent_sessions", [], |row| row.get(0))?;
let total_tokens: i64 = conn.query_row(
"SELECT COALESCE(SUM(COALESCE(accumulated_total_tokens, total_tokens, 0)), 0)
FROM agent_sessions",
[],
|row| row.get(0),
)?;
Ok(SessionInsights {
total_sessions: total_sessions as usize,
total_tokens: 0,
total_tokens,
})
}
@@ -528,6 +658,36 @@ impl SessionStore for LimeSessionStore {
)
.await?;
self.update_session_name(&new_session.id, session.name.clone(), session.user_set_name)
.await?;
self.update_extension_data(&new_session.id, session.extension_data.clone())
.await?;
self.update_token_stats(
&new_session.id,
TokenStatsUpdate {
schedule_id: session.schedule_id.clone(),
total_tokens: session.total_tokens,
input_tokens: session.input_tokens,
output_tokens: session.output_tokens,
accumulated_total: session.accumulated_total_tokens,
accumulated_input: session.accumulated_input_tokens,
accumulated_output: session.accumulated_output_tokens,
},
)
.await?;
self.update_provider_config(
&new_session.id,
session.provider_name.clone(),
session.model_config.clone(),
)
.await?;
self.update_recipe(
&new_session.id,
session.recipe.clone(),
session.user_recipe_values.clone(),
)
.await?;
if let Some(conversation) = &session.conversation {
self.replace_conversation(&new_session.id, conversation)
.await?;
@@ -538,15 +698,47 @@ impl SessionStore for LimeSessionStore {
async fn copy_session(&self, session_id: &str, new_name: String) -> Result<Session> {
let original = self.get_session(session_id, true).await?;
let created_session_name = new_name.clone();
let persisted_session_name = new_name.clone();
let new_session = self
.create_session(
original.working_dir.clone(),
new_name,
created_session_name,
original.session_type,
)
.await?;
self.update_session_name(&new_session.id, persisted_session_name, true)
.await?;
self.update_extension_data(&new_session.id, original.extension_data.clone())
.await?;
self.update_token_stats(
&new_session.id,
TokenStatsUpdate {
schedule_id: original.schedule_id.clone(),
total_tokens: original.total_tokens,
input_tokens: original.input_tokens,
output_tokens: original.output_tokens,
accumulated_total: original.accumulated_total_tokens,
accumulated_input: original.accumulated_input_tokens,
accumulated_output: original.accumulated_output_tokens,
},
)
.await?;
self.update_provider_config(
&new_session.id,
original.provider_name.clone(),
original.model_config.clone(),
)
.await?;
self.update_recipe(
&new_session.id,
original.recipe.clone(),
original.user_recipe_values.clone(),
)
.await?;
if let Some(conversation) = &original.conversation {
self.replace_conversation(&new_session.id, conversation)
.await?;
@@ -573,25 +765,59 @@ impl SessionStore for LimeSessionStore {
&self,
session_id: &str,
name: String,
_user_set: bool,
user_set: bool,
) -> Result<()> {
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let now = Utc::now().to_rfc3339();
conn.execute(
"UPDATE agent_sessions SET title = ? WHERE id = ?",
rusqlite::params![name, session_id],
"UPDATE agent_sessions SET title = ?1, user_set_name = ?2, updated_at = ?3 WHERE id = ?4",
rusqlite::params![name, user_set, now, session_id],
)?;
Ok(())
}
async fn update_extension_data(
&self,
_session_id: &str,
_extension_data: ExtensionData,
session_id: &str,
extension_data: ExtensionData,
) -> Result<()> {
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let now = Utc::now().to_rfc3339();
let extension_data_json = serde_json::to_string(&extension_data)
.map_err(|e| anyhow!("序列化 extension_data 失败: {e}"))?;
conn.execute(
"UPDATE agent_sessions SET extension_data_json = ?1, updated_at = ?2 WHERE id = ?3",
rusqlite::params![extension_data_json, now, session_id],
)?;
Ok(())
}
async fn update_token_stats(&self, _session_id: &str, _stats: TokenStatsUpdate) -> Result<()> {
async fn update_token_stats(&self, session_id: &str, stats: TokenStatsUpdate) -> Result<()> {
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let now = Utc::now().to_rfc3339();
conn.execute(
"UPDATE agent_sessions SET
total_tokens = COALESCE(?1, total_tokens),
input_tokens = COALESCE(?2, input_tokens),
output_tokens = COALESCE(?3, output_tokens),
accumulated_total_tokens = COALESCE(?4, accumulated_total_tokens),
accumulated_input_tokens = COALESCE(?5, accumulated_input_tokens),
accumulated_output_tokens = COALESCE(?6, accumulated_output_tokens),
schedule_id = COALESCE(?7, schedule_id),
updated_at = ?8
WHERE id = ?9",
rusqlite::params![
stats.total_tokens,
stats.input_tokens,
stats.output_tokens,
stats.accumulated_total,
stats.accumulated_input,
stats.accumulated_output,
Self::normalize_optional_text(stats.schedule_id),
now,
session_id,
],
)?;
Ok(())
}
@@ -601,27 +827,67 @@ impl SessionStore for LimeSessionStore {
provider_name: Option<String>,
model_config: Option<ModelConfig>,
) -> Result<()> {
if let Some(model_name) = model_config
let normalized_provider_name = Self::normalize_optional_text(provider_name);
let normalized_model_name = model_config
.as_ref()
.map(|config| config.model_name.trim().to_string())
.filter(|value| !value.is_empty())
.or(provider_name.filter(|value| !value.trim().is_empty()))
{
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
conn.execute(
"UPDATE agent_sessions SET model = ? WHERE id = ?",
rusqlite::params![model_name, session_id],
)?;
.filter(|value| !value.is_empty());
let model_config_json = model_config
.as_ref()
.map(serde_json::to_string)
.transpose()
.map_err(|e| anyhow!("序列化 model_config 失败: {e}"))?;
if normalized_provider_name.is_none() && normalized_model_name.is_none() {
return Ok(());
}
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let now = Utc::now().to_rfc3339();
conn.execute(
"UPDATE agent_sessions SET
provider_name = COALESCE(?1, provider_name),
model = COALESCE(?2, model),
model_config_json = CASE WHEN ?3 IS NULL THEN model_config_json ELSE ?3 END,
updated_at = ?4
WHERE id = ?5",
rusqlite::params![
normalized_provider_name,
normalized_model_name,
model_config_json,
now,
session_id,
],
)?;
Ok(())
}
async fn update_recipe(
&self,
_session_id: &str,
_recipe: Option<Recipe>,
_user_recipe_values: Option<HashMap<String, String>>,
session_id: &str,
recipe: Option<Recipe>,
user_recipe_values: Option<HashMap<String, String>>,
) -> Result<()> {
let conn = self.db.lock().map_err(|e| anyhow!("数据库锁定失败: {e}"))?;
let now = Utc::now().to_rfc3339();
let recipe_json = recipe
.as_ref()
.map(serde_json::to_string)
.transpose()
.map_err(|e| anyhow!("序列化 recipe 失败: {e}"))?;
let user_recipe_values_json = user_recipe_values
.as_ref()
.map(serde_json::to_string)
.transpose()
.map_err(|e| anyhow!("序列化 user_recipe_values 失败: {e}"))?;
conn.execute(
"UPDATE agent_sessions SET
recipe_json = ?1,
user_recipe_values_json = ?2,
updated_at = ?3
WHERE id = ?4",
rusqlite::params![recipe_json, user_recipe_values_json, now, session_id],
)?;
Ok(())
}
@@ -788,7 +1054,41 @@ mod tests {
use aster::session::{SessionStore, SessionType};
use lime_core::database::schema::create_tables;
use rusqlite::Connection;
use std::ffi::OsString;
use std::sync::{Arc, Mutex};
use tempfile::tempdir;
fn env_lock() -> &'static Mutex<()> {
static LOCK: std::sync::OnceLock<Mutex<()>> = std::sync::OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
}
struct EnvGuard {
values: Vec<(&'static str, Option<OsString>)>,
}
impl EnvGuard {
fn set(entries: &[(&'static str, OsString)]) -> Self {
let mut values = Vec::new();
for (key, value) in entries {
values.push((*key, std::env::var_os(key)));
std::env::set_var(key, value);
}
Self { values }
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
for (key, previous) in self.values.drain(..) {
if let Some(value) = previous {
std::env::set_var(key, value);
} else {
std::env::remove_var(key);
}
}
}
}
fn setup_test_store() -> LimeSessionStore {
let conn = Connection::open_in_memory().expect("创建内存数据库失败");
@@ -828,4 +1128,175 @@ mod tests {
assert_eq!(persisted_model, "gpt-4.1");
}
#[tokio::test]
async fn get_session_should_prefer_default_workspace_root_when_missing_row() {
let store = setup_test_store();
let workspace_root = std::env::temp_dir().join("lime-aster-default-workspace");
let conn = store.db.lock().expect("锁数据库");
conn.execute(
"INSERT INTO workspaces (id, name, workspace_type, root_path, is_default, settings_json, created_at, updated_at)
VALUES (?1, ?2, ?3, ?4, 1, '{}', 0, 0)",
rusqlite::params![
"workspace-default",
"默认工作区",
"general",
workspace_root.to_string_lossy().to_string(),
],
)
.expect("插入默认 workspace 失败");
drop(conn);
let session = store
.get_session("missing-default-workspace-session", false)
.await
.expect("读取缺失会话失败");
assert_eq!(session.working_dir, workspace_root);
}
#[tokio::test]
async fn get_session_should_fallback_to_app_paths_default_project_dir() {
let _env_guard = env_lock().lock().expect("锁环境变量");
let temp = tempdir().expect("创建临时目录失败");
let home = temp.path().join("home");
let app_data = temp.path().join("appdata");
std::fs::create_dir_all(&home).expect("创建 home 目录失败");
std::fs::create_dir_all(&app_data).expect("创建 appdata 目录失败");
let _guard = EnvGuard::set(&[
("HOME", home.as_os_str().to_os_string()),
("XDG_DATA_HOME", app_data.as_os_str().to_os_string()),
("APPDATA", app_data.as_os_str().to_os_string()),
("LOCALAPPDATA", app_data.as_os_str().to_os_string()),
]);
let store = setup_test_store();
let session = store
.get_session("missing-fallback-session", false)
.await
.expect("读取缺失会话失败");
let expected = app_paths::resolve_default_project_dir().expect("解析默认项目目录失败");
assert_eq!(session.working_dir, expected);
assert!(session.working_dir.is_absolute());
assert!(session
.working_dir
.ends_with(PathBuf::from("projects").join("default")));
assert!(!session
.working_dir
.to_string_lossy()
.contains(".lime/projects/default"));
}
#[tokio::test]
async fn update_session_metadata_should_roundtrip() {
let store = setup_test_store();
let session = store
.create_session(
PathBuf::from("."),
"元数据测试".to_string(),
SessionType::SubAgent,
)
.await
.expect("创建会话失败");
let mut extension_data = ExtensionData::new();
extension_data.set_extension_state("todo", "v0", serde_json::json!({"items":["a"]}));
store
.update_session_name(&session.id, "已命名会话".to_string(), true)
.await
.expect("更新名称失败");
store
.update_extension_data(&session.id, extension_data.clone())
.await
.expect("更新 extension_data 失败");
store
.update_token_stats(
&session.id,
TokenStatsUpdate {
schedule_id: Some("job-1".to_string()),
total_tokens: Some(100),
input_tokens: Some(60),
output_tokens: Some(40),
accumulated_total: Some(300),
accumulated_input: Some(180),
accumulated_output: Some(120),
},
)
.await
.expect("更新 token 统计失败");
store
.update_provider_config(
&session.id,
Some("openai".to_string()),
Some(ModelConfig::new("gpt-4.1").expect("model config")),
)
.await
.expect("更新 provider 配置失败");
store
.update_recipe(
&session.id,
Some(Recipe {
version: "1.0.0".to_string(),
title: "demo".to_string(),
description: "demo recipe".to_string(),
instructions: None,
prompt: None,
extensions: None,
settings: None,
activities: None,
author: None,
parameters: None,
response: None,
sub_recipes: None,
retry: None,
}),
Some(HashMap::from([(
"temperature".to_string(),
"0.2".to_string(),
)])),
)
.await
.expect("更新 recipe 失败");
let loaded = store
.get_session(&session.id, false)
.await
.expect("读取会话失败");
assert_eq!(loaded.name, "已命名会话");
assert!(loaded.user_set_name);
assert_eq!(loaded.session_type, SessionType::SubAgent);
assert_eq!(loaded.total_tokens, Some(100));
assert_eq!(loaded.accumulated_total_tokens, Some(300));
assert_eq!(loaded.schedule_id.as_deref(), Some("job-1"));
assert_eq!(loaded.provider_name.as_deref(), Some("openai"));
assert_eq!(
loaded
.model_config
.as_ref()
.map(|config| config.model_name.as_str()),
Some("gpt-4.1")
);
assert_eq!(
loaded
.extension_data
.get_extension_state("todo", "v0")
.cloned(),
extension_data.get_extension_state("todo", "v0").cloned()
);
assert_eq!(
loaded.recipe.as_ref().map(|recipe| recipe.title.as_str()),
Some("demo")
);
assert_eq!(
loaded
.user_recipe_values
.as_ref()
.and_then(|values| values.get("temperature"))
.map(String::as_str),
Some("0.2")
);
}
}
src-tauri/crates/services/src/lib.rs
-2
View File
@@ -43,7 +43,6 @@
//! - `session_context_service` - 会话上下文服务
//! - `ai_summary_service` - AI 摘要服务
//! - `project_context_builder` - 项目上下文构建器
//! - `tool_hooks_service` - 工具钩子服务
//! - `kiro_event_service` - Kiro 事件服务
//! - `api_key_provider_service` - API Key Provider 服务
//! - `provider_pool_service` - Provider 池服务
@@ -90,7 +89,6 @@ pub mod content_creator;
pub mod ai_summary_service;
pub mod project_context_builder;
pub mod session_context_service;
pub mod tool_hooks_service;
// 事件服务
pub mod kiro_event_service;
src-tauri/crates/services/src/session_context_service.rs
+20 -153
View File
@@ -4,7 +4,6 @@
use crate::ai_summary_service::AISummaryService;
use lime_core::database::dao::chat::{ChatDao, ChatMessage as UnifiedChatMessage, ChatMode};
use lime_core::database::load_pending_general_session_messages;
use lime_core::general_chat::{ChatMessage, MessageRole};
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
@@ -184,52 +183,27 @@ impl SessionContextService {
let session =
ChatDao::get_session(conn, session_id).map_err(|e| format!("获取统一会话失败: {e}"))?;
if let Some(session) = session {
if session.mode != ChatMode::General {
debug!(
"会话 {} 不是 general 模式，跳过上下文加载: {:?}",
session_id, session.mode
);
return Ok(vec![]);
}
let Some(session) = session else {
debug!("会话 {} 未命中 unified chat，跳过上下文加载", session_id);
return Ok(vec![]);
};
let unified_messages = ChatDao::get_messages(conn, session_id, None)
.map_err(|e| format!("获取统一消息失败: {e}"))?;
return Ok(unified_messages
.into_iter()
.map(Self::convert_unified_message)
.collect());
if session.mode != ChatMode::General {
debug!(
"会话 {} 不是 general 模式，跳过上下文加载: {:?}",
session_id, session.mode
);
return Ok(vec![]);
}
debug!(
"会话 {} 未命中 unified chat，回退待迁移 general 历史表",
session_id
);
Self::load_pending_general_messages(conn, session_id)
}
fn load_pending_general_messages(
conn: &Connection,
session_id: &str,
) -> Result<Vec<ChatMessage>, String> {
let messages = load_pending_general_session_messages(conn, session_id)
.map_err(|e| format!("查询待迁移 general 消息失败: {e}"))?;
Ok(messages
.into_iter()
.map(|message| ChatMessage {
id: message.id,
session_id: message.session_id,
role: Self::convert_legacy_role(&message.role),
content: message.content,
blocks: None,
status: "complete".to_string(),
created_at: message.created_at,
metadata: None,
ChatDao::get_messages(conn, session_id, None)
.map_err(|e| format!("获取统一消息失败: {e}"))
.map(|messages| {
messages
.into_iter()
.map(Self::convert_unified_message)
.collect()
})
.collect())
}
fn convert_unified_message(message: UnifiedChatMessage) -> ChatMessage {
@@ -253,14 +227,6 @@ impl SessionContextService {
}
}
fn convert_legacy_role(role: &str) -> MessageRole {
match role {
"assistant" => MessageRole::Assistant,
"system" => MessageRole::System,
_ => MessageRole::User,
}
}
fn extract_unified_text_content(content: &serde_json::Value) -> String {
match content {
serde_json::Value::Null => String::new(),
@@ -604,7 +570,6 @@ mod tests {
use lime_core::database::dao::chat::{
ChatDao, ChatMessage as UnifiedChatMessage, ChatMode, ChatSession as UnifiedChatSession,
};
use lime_core::general_chat::ChatSession;
use rusqlite::Connection;
fn setup_test_db() -> Connection {
@@ -649,36 +614,6 @@ mod tests {
)
.unwrap();
// 创建会话表
conn.execute(
"CREATE TABLE general_chat_sessions (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL,
metadata TEXT
)",
[],
)
.unwrap();
// 创建消息表
conn.execute(
"CREATE TABLE general_chat_messages (
id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
role TEXT NOT NULL CHECK (role IN ('user', 'assistant', 'system')),
content TEXT NOT NULL,
blocks TEXT,
status TEXT NOT NULL DEFAULT 'complete',
created_at INTEGER NOT NULL,
metadata TEXT,
FOREIGN KEY (session_id) REFERENCES general_chat_sessions(id) ON DELETE CASCADE
)",
[],
)
.unwrap();
conn.execute("PRAGMA foreign_keys = ON", []).unwrap();
conn
}
@@ -749,45 +684,6 @@ mod tests {
.collect()
}
fn insert_legacy_session(conn: &Connection, session: &ChatSession) {
conn.execute(
"INSERT INTO general_chat_sessions (id, name, created_at, updated_at, metadata)
VALUES (?1, ?2, ?3, ?4, ?5)",
rusqlite::params![
session.id,
session.name,
session.created_at,
session.updated_at,
Option::<String>::None,
],
)
.unwrap();
}
fn insert_legacy_message(conn: &Connection, message: &ChatMessage) {
let role = match message.role {
MessageRole::User => "user",
MessageRole::Assistant => "assistant",
MessageRole::System => "system",
};
conn.execute(
"INSERT INTO general_chat_messages (id, session_id, role, content, blocks, status, created_at, metadata)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
rusqlite::params![
message.id,
message.session_id,
role,
message.content,
Option::<String>::None,
message.status,
message.created_at,
Option::<String>::None,
],
)
.unwrap();
}
#[tokio::test]
async fn test_context_service_creation() {
let conn = Arc::new(Mutex::new(setup_test_db()));
@@ -880,47 +776,22 @@ mod tests {
}
#[tokio::test]
async fn test_get_effective_context_falls_back_to_legacy_messages() {
async fn test_get_effective_context_returns_empty_when_unified_session_missing() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let service = SessionContextService::new(conn.clone(), ContextWindowConfig::default());
let legacy_session = ChatSession {
id: "legacy-session".to_string(),
name: "旧会话".to_string(),
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
metadata: None,
};
{
let conn_guard = conn.lock().unwrap();
insert_legacy_session(&conn_guard, &legacy_session);
for message in create_test_messages("legacy-session", 3) {
insert_legacy_message(&conn_guard, &message);
}
}
let context = service
.get_effective_context("legacy-session")
.await
.unwrap();
assert_eq!(context.len(), 3);
assert_eq!(context[0].content, "这是第 1 条消息，包含一些测试内容");
assert!(context.is_empty());
}
#[tokio::test]
async fn test_get_effective_context_skips_legacy_fallback_after_general_migration_completed() {
async fn test_get_effective_context_keeps_empty_when_general_migration_completed() {
let conn = Arc::new(Mutex::new(setup_test_db()));
let service = SessionContextService::new(conn.clone(), ContextWindowConfig::default());
let legacy_session = ChatSession {
id: "legacy-session".to_string(),
name: "旧会话".to_string(),
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
metadata: None,
};
{
let conn_guard = conn.lock().unwrap();
conn_guard
@@ -929,10 +800,6 @@ mod tests {
rusqlite::params!["migrated_general_chat_to_unified", "true"],
)
.unwrap();
insert_legacy_session(&conn_guard, &legacy_session);
for message in create_test_messages("legacy-session", 3) {
insert_legacy_message(&conn_guard, &message);
}
}
let context = service
src-tauri/crates/services/src/tool_hooks_service.rs
-715
View File
@@ -1,715 +0,0 @@
//! 工具钩子管理服务
//!
//! 提供工具执行前后的钩子机制，用于自动化上下文记忆管理
use crate::context_memory_service::{ContextMemoryService, MemoryEntry, MemoryFileType};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tracing::{debug, error, info};
/// 钩子触发时机
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HookTrigger {
/// 会话开始时
SessionStart,
/// 工具使用前
PreToolUse,
/// 工具使用后
PostToolUse,
/// 会话停止时
Stop,
}
/// 钩子动作类型
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HookAction {
/// 保存发现
SaveFinding {
title: String,
content: String,
tags: Vec<String>,
priority: u8,
},
/// 更新任务计划
UpdateTaskPlan {
title: String,
content: String,
priority: u8,
},
/// 记录进度
LogProgress { title: String, content: String },
/// 记录错误
RecordError {
error_description: String,
attempted_solution: String,
},
/// 自定义动作
Custom {
action_type: String,
parameters: HashMap<String, String>,
},
}
/// 钩子规则
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HookRule {
/// 规则 ID
pub id: String,
/// 规则名称
pub name: String,
/// 描述
pub description: String,
/// 触发时机
pub trigger: HookTrigger,
/// 触发条件
pub conditions: Vec<HookCondition>,
/// 执行动作
pub actions: Vec<HookAction>,
/// 是否启用
pub enabled: bool,
/// 优先级 (数字越小优先级越高)
pub priority: u32,
/// 创建时间
pub created_at: i64,
}
/// 钩子条件
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HookCondition {
/// 工具名称匹配
ToolNameEquals(String),
/// 工具名称包含
ToolNameContains(String),
/// 消息内容包含
MessageContains(String),
/// 会话消息数量大于
MessageCountGreaterThan(usize),
/// 错误次数大于
ErrorCountGreaterThan(u32),
/// 自定义条件
Custom {
condition_type: String,
parameters: HashMap<String, String>,
},
}
/// 钩子执行上下文
#[derive(Debug, Clone)]
pub struct HookContext {
/// 会话 ID
pub session_id: String,
/// 工具名称（如果适用）
pub tool_name: Option<String>,
/// 工具参数（如果适用）
pub tool_parameters: Option<HashMap<String, String>>,
/// 工具结果（如果适用）
pub tool_result: Option<String>,
/// 消息内容
pub message_content: Option<String>,
/// 会话消息数量
pub message_count: usize,
/// 错误信息（如果适用）
pub error_info: Option<String>,
/// 额外元数据
pub metadata: HashMap<String, String>,
}
/// 工具钩子管理器
pub struct ToolHooksService {
/// 钩子规则
rules: Arc<Mutex<Vec<HookRule>>>,
/// 上下文记忆服务
memory_service: Arc<ContextMemoryService>,
/// 执行统计
execution_stats: Arc<Mutex<HashMap<String, HookExecutionStats>>>,
}
/// 钩子执行统计
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct HookExecutionStats {
/// 执行次数
pub execution_count: u64,
/// 成功次数
pub success_count: u64,
/// 失败次数
pub failure_count: u64,
/// 最后执行时间
pub last_execution_at: i64,
/// 平均执行时间（毫秒）
pub average_execution_time_ms: f64,
}
impl ToolHooksService {
/// 创建新的工具钩子服务
pub fn new(memory_service: Arc<ContextMemoryService>) -> Self {
let service = Self {
rules: Arc::new(Mutex::new(Vec::new())),
memory_service,
execution_stats: Arc::new(Mutex::new(HashMap::new())),
};
// 注册默认钩子规则
service.register_default_hooks();
service
}
/// 注册默认钩子规则
fn register_default_hooks(&self) {
let default_rules = vec![
// 会话开始时创建任务计划
HookRule {
id: "session-start-task-plan".to_string(),
name: "会话开始任务计划".to_string(),
description: "会话开始时自动创建任务计划记录".to_string(),
trigger: HookTrigger::SessionStart,
conditions: vec![],
actions: vec![HookAction::UpdateTaskPlan {
title: "会话任务计划".to_string(),
content: "新会话开始，等待用户输入任务目标...".to_string(),
priority: 3,
}],
enabled: true,
priority: 1,
created_at: chrono::Utc::now().timestamp_millis(),
},
// 工具使用后记录进度（2-Action 规则）
HookRule {
id: "post-tool-progress-log".to_string(),
name: "工具使用进度记录".to_string(),
description: "工具使用后自动记录进度（2-Action 规则）".to_string(),
trigger: HookTrigger::PostToolUse,
conditions: vec![],
actions: vec![HookAction::LogProgress {
title: "工具执行记录".to_string(),
content: "工具执行完成，结果已记录".to_string(),
}],
enabled: true,
priority: 2,
created_at: chrono::Utc::now().timestamp_millis(),
},
// 重要发现自动保存
HookRule {
id: "important-finding-save".to_string(),
name: "重要发现保存".to_string(),
description: "检测到重要信息时自动保存".to_string(),
trigger: HookTrigger::PostToolUse,
conditions: vec![
HookCondition::MessageContains("重要".to_string()),
HookCondition::MessageContains("发现".to_string()),
],
actions: vec![HookAction::SaveFinding {
title: "重要发现".to_string(),
content: "检测到重要信息，已自动保存".to_string(),
tags: vec!["重要".to_string(), "自动保存".to_string()],
priority: 4,
}],
enabled: true,
priority: 3,
created_at: chrono::Utc::now().timestamp_millis(),
},
// 错误自动记录
HookRule {
id: "error-auto-record".to_string(),
name: "错误自动记录".to_string(),
description: "检测到错误时自动记录".to_string(),
trigger: HookTrigger::PostToolUse,
conditions: vec![HookCondition::MessageContains("错误".to_string())],
actions: vec![HookAction::RecordError {
error_description: "检测到错误".to_string(),
attempted_solution: "正在尝试解决".to_string(),
}],
enabled: true,
priority: 1,
created_at: chrono::Utc::now().timestamp_millis(),
},
// 会话停止时保存摘要
HookRule {
id: "session-stop-summary".to_string(),
name: "会话停止摘要".to_string(),
description: "会话停止时保存会话摘要".to_string(),
trigger: HookTrigger::Stop,
conditions: vec![HookCondition::MessageCountGreaterThan(5)],
actions: vec![HookAction::SaveFinding {
title: "会话摘要".to_string(),
content: "会话已结束，主要成果和发现已记录".to_string(),
tags: vec!["摘要".to_string(), "会话结束".to_string()],
priority: 3,
}],
enabled: true,
priority: 2,
created_at: chrono::Utc::now().timestamp_millis(),
},
];
let mut rules = self.rules.lock().unwrap();
rules.extend(default_rules);
info!("已注册 {} 个默认钩子规则", rules.len());
}
/// 执行钩子
pub fn execute_hooks(&self, trigger: HookTrigger, context: &HookContext) -> Result<(), String> {
let rules = self.rules.lock().map_err(|e| e.to_string())?;
// 获取匹配的规则并按优先级排序
let mut matching_rules: Vec<_> = rules
.iter()
.filter(|rule| rule.enabled && rule.trigger == trigger)
.filter(|rule| self.evaluate_conditions(rule, context))
.collect();
matching_rules.sort_by_key(|rule| rule.priority);
debug!(
"触发钩子 {:?}，匹配到 {} 个规则 (会话: {})",
trigger,
matching_rules.len(),
context.session_id
);
// 执行匹配的规则
for rule in matching_rules {
if let Err(e) = self.execute_rule(rule, context) {
error!("执行钩子规则失败 {}: {}", rule.name, e);
self.update_execution_stats(&rule.id, false, 0.0);
} else {
self.update_execution_stats(&rule.id, true, 0.0);
}
}
Ok(())
}
/// 评估钩子条件
fn evaluate_conditions(&self, rule: &HookRule, context: &HookContext) -> bool {
if rule.conditions.is_empty() {
return true;
}
for condition in &rule.conditions {
if !self.evaluate_single_condition(condition, context) {
return false;
}
}
true
}
/// 评估单个条件
fn evaluate_single_condition(&self, condition: &HookCondition, context: &HookContext) -> bool {
match condition {
HookCondition::ToolNameEquals(name) => context.tool_name.as_ref() == Some(name),
HookCondition::ToolNameContains(substring) => context
.tool_name
.as_ref()
.is_some_and(|tn| tn.contains(substring)),
HookCondition::MessageContains(substring) => {
context
.message_content
.as_ref()
.is_some_and(|mc| mc.contains(substring))
|| context
.tool_result
.as_ref()
.is_some_and(|tr| tr.contains(substring))
}
HookCondition::MessageCountGreaterThan(count) => context.message_count > *count,
HookCondition::ErrorCountGreaterThan(_count) => {
// 这里可以从 memory_service 获取错误计数
context.error_info.is_some()
}
HookCondition::Custom {
condition_type: _,
parameters: _,
} => {
// 自定义条件的实现
true
}
}
}
/// 执行钩子规则
fn execute_rule(&self, rule: &HookRule, context: &HookContext) -> Result<(), String> {
let start_time = std::time::Instant::now();
for action in &rule.actions {
self.execute_action(action, context)?;
}
let execution_time = start_time.elapsed().as_millis() as f64;
self.update_execution_stats(&rule.id, true, execution_time);
debug!(
"执行钩子规则成功: {} (耗时: {:.2}ms)",
rule.name, execution_time
);
Ok(())
}
/// 执行钩子动作
fn execute_action(&self, action: &HookAction, context: &HookContext) -> Result<(), String> {
match action {
HookAction::SaveFinding {
title,
content,
tags,
priority,
} => {
let entry = MemoryEntry {
id: uuid::Uuid::new_v4().to_string(),
session_id: context.session_id.clone(),
file_type: MemoryFileType::Findings,
title: self.interpolate_template(title, context),
content: self.interpolate_template(content, context),
tags: tags.clone(),
priority: *priority,
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
archived: false,
};
self.memory_service.save_memory_entry(&entry)?;
}
HookAction::UpdateTaskPlan {
title,
content,
priority,
} => {
let entry = MemoryEntry {
id: uuid::Uuid::new_v4().to_string(),
session_id: context.session_id.clone(),
file_type: MemoryFileType::TaskPlan,
title: self.interpolate_template(title, context),
content: self.interpolate_template(content, context),
tags: vec!["任务计划".to_string()],
priority: *priority,
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
archived: false,
};
self.memory_service.save_memory_entry(&entry)?;
}
HookAction::LogProgress { title, content } => {
let entry = MemoryEntry {
id: uuid::Uuid::new_v4().to_string(),
session_id: context.session_id.clone(),
file_type: MemoryFileType::Progress,
title: self.interpolate_template(title, context),
content: self.interpolate_template(content, context),
tags: vec!["进度".to_string()],
priority: 2,
created_at: chrono::Utc::now().timestamp_millis(),
updated_at: chrono::Utc::now().timestamp_millis(),
archived: false,
};
self.memory_service.save_memory_entry(&entry)?;
}
HookAction::RecordError {
error_description,
attempted_solution,
} => {
let error_desc = self.interpolate_template(error_description, context);
let solution = self.interpolate_template(attempted_solution, context);
self.memory_service
.record_error(&context.session_id, &error_desc, &solution)?;
}
HookAction::Custom {
action_type: _,
parameters: _,
} => {
// 自定义动作的实现
debug!("执行自定义钩子动作");
}
}
Ok(())
}
/// 模板插值
fn interpolate_template(&self, template: &str, context: &HookContext) -> String {
let mut result = template.to_string();
// 替换常见的模板变量
result = result.replace("{session_id}", &context.session_id);
if let Some(tool_name) = &context.tool_name {
result = result.replace("{tool_name}", tool_name);
}
if let Some(message_content) = &context.message_content {
let preview = if message_content.len() > 100 {
format!("{}...", &message_content[..100])
} else {
message_content.clone()
};
result = result.replace("{message_preview}", &preview);
}
if let Some(tool_result) = &context.tool_result {
let preview = if tool_result.len() > 200 {
format!("{}...", &tool_result[..200])
} else {
tool_result.clone()
};
result = result.replace("{tool_result_preview}", &preview);
}
result = result.replace("{message_count}", &context.message_count.to_string());
result = result.replace(
"{timestamp}",
&chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string(),
);
// 替换元数据变量
for (key, value) in &context.metadata {
result = result.replace(&format!("{{{key}}}"), value);
}
result
}
/// 更新执行统计
fn update_execution_stats(&self, rule_id: &str, success: bool, execution_time_ms: f64) {
let mut stats = self.execution_stats.lock().unwrap();
let entry = stats.entry(rule_id.to_string()).or_default();
entry.execution_count += 1;
if success {
entry.success_count += 1;
} else {
entry.failure_count += 1;
}
entry.last_execution_at = chrono::Utc::now().timestamp_millis();
// 更新平均执行时间
if execution_time_ms > 0.0 {
let total_time = entry.average_execution_time_ms * (entry.execution_count - 1) as f64;
entry.average_execution_time_ms =
(total_time + execution_time_ms) / entry.execution_count as f64;
}
}
/// 添加钩子规则
pub fn add_hook_rule(&self, rule: HookRule) -> Result<(), String> {
let mut rules = self.rules.lock().map_err(|e| e.to_string())?;
// 检查是否已存在相同 ID 的规则
if rules.iter().any(|r| r.id == rule.id) {
return Err(format!("钩子规则 ID 已存在: {}", rule.id));
}
rules.push(rule.clone());
info!("已添加钩子规则: {}", rule.name);
Ok(())
}
/// 移除钩子规则
pub fn remove_hook_rule(&self, rule_id: &str) -> Result<(), String> {
let mut rules = self.rules.lock().map_err(|e| e.to_string())?;
let initial_len = rules.len();
rules.retain(|r| r.id != rule_id);
if rules.len() == initial_len {
return Err(format!("未找到钩子规则: {rule_id}"));
}
info!("已移除钩子规则: {}", rule_id);
Ok(())
}
/// 启用/禁用钩子规则
pub fn toggle_hook_rule(&self, rule_id: &str, enabled: bool) -> Result<(), String> {
let mut rules = self.rules.lock().map_err(|e| e.to_string())?;
if let Some(rule) = rules.iter_mut().find(|r| r.id == rule_id) {
rule.enabled = enabled;
info!(
"钩子规则 {} 已{}",
rule.name,
if enabled { "启用" } else { "禁用" }
);
Ok(())
} else {
Err(format!("未找到钩子规则: {rule_id}"))
}
}
/// 获取所有钩子规则
pub fn get_hook_rules(&self) -> Result<Vec<HookRule>, String> {
let rules = self.rules.lock().map_err(|e| e.to_string())?;
Ok(rules.clone())
}
/// 获取执行统计
pub fn get_execution_stats(&self) -> Result<HashMap<String, HookExecutionStats>, String> {
let stats = self.execution_stats.lock().map_err(|e| e.to_string())?;
Ok(stats.clone())
}
/// 清理执行统计
pub fn clear_execution_stats(&self) -> Result<(), String> {
let mut stats = self.execution_stats.lock().map_err(|e| e.to_string())?;
stats.clear();
info!("已清理钩子执行统计");
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context_memory_service::ContextMemoryConfig;
use tempfile::TempDir;
fn create_test_services() -> (Arc<ContextMemoryService>, ToolHooksService, TempDir) {
let temp_dir = TempDir::new().unwrap();
let memory_config = ContextMemoryConfig {
memory_dir: temp_dir.path().to_path_buf(),
..Default::default()
};
let memory_service = Arc::new(ContextMemoryService::new(memory_config).unwrap());
let hooks_service = ToolHooksService::new(memory_service.clone());
(memory_service, hooks_service, temp_dir)
}
#[test]
fn test_hooks_service_creation() {
let (_memory_service, hooks_service, _temp_dir) = create_test_services();
let rules = hooks_service.get_hook_rules().unwrap();
assert!(!rules.is_empty()); // 应该有默认规则
}
#[test]
fn test_session_start_hook() {
let (_memory_service, hooks_service, _temp_dir) = create_test_services();
let context = HookContext {
session_id: "test-session".to_string(),
tool_name: None,
tool_parameters: None,
tool_result: None,
message_content: None,
message_count: 0,
error_info: None,
metadata: HashMap::new(),
};
hooks_service
.execute_hooks(HookTrigger::SessionStart, &context)
.unwrap();
// 验证任务计划是否被创建
let memories = _memory_service
.get_session_memories("test-session", Some(MemoryFileType::TaskPlan))
.unwrap();
assert!(!memories.is_empty());
}
#[test]
fn test_error_recording_hook() {
let (_memory_service, hooks_service, _temp_dir) = create_test_services();
let context = HookContext {
session_id: "test-session".to_string(),
tool_name: Some("test_tool".to_string()),
tool_parameters: None,
tool_result: Some("发生了一个错误".to_string()),
message_content: Some("这里有一个错误需要处理".to_string()),
message_count: 5,
error_info: Some("测试错误".to_string()),
metadata: HashMap::new(),
};
hooks_service
.execute_hooks(HookTrigger::PostToolUse, &context)
.unwrap();
// 验证错误是否被记录
let stats = _memory_service.get_memory_stats("test-session").unwrap();
assert!(stats.unresolved_errors > 0);
}
#[test]
fn test_custom_hook_rule() {
let (_memory_service, hooks_service, _temp_dir) = create_test_services();
let custom_rule = HookRule {
id: "custom-test-rule".to_string(),
name: "自定义测试规则".to_string(),
description: "测试自定义钩子规则".to_string(),
trigger: HookTrigger::PostToolUse,
conditions: vec![HookCondition::ToolNameEquals("custom_tool".to_string())],
actions: vec![HookAction::SaveFinding {
title: "自定义发现".to_string(),
content: "这是一个自定义钩子触发的发现".to_string(),
tags: vec!["自定义".to_string()],
priority: 3,
}],
enabled: true,
priority: 1,
created_at: chrono::Utc::now().timestamp_millis(),
};
hooks_service.add_hook_rule(custom_rule).unwrap();
let context = HookContext {
session_id: "test-session".to_string(),
tool_name: Some("custom_tool".to_string()),
tool_parameters: None,
tool_result: None,
message_content: None,
message_count: 0,
error_info: None,
metadata: HashMap::new(),
};
hooks_service
.execute_hooks(HookTrigger::PostToolUse, &context)
.unwrap();
// 验证自定义发现是否被保存
let memories = _memory_service
.get_session_memories("test-session", Some(MemoryFileType::Findings))
.unwrap();
assert!(memories.iter().any(|m| m.title == "自定义发现"));
}
#[test]
fn test_template_interpolation() {
let (_memory_service, hooks_service, _temp_dir) = create_test_services();
let context = HookContext {
session_id: "test-session-123".to_string(),
tool_name: Some("test_tool".to_string()),
tool_parameters: None,
tool_result: None,
message_content: Some("这是测试消息".to_string()),
message_count: 42,
error_info: None,
metadata: {
let mut map = HashMap::new();
map.insert("custom_var".to_string(), "custom_value".to_string());
map
},
};
let template = "会话 {session_id} 使用工具 {tool_name}，消息数量: {message_count}，自定义变量: {custom_var}";
let result = hooks_service.interpolate_template(template, &context);
assert!(result.contains("test-session-123"));
assert!(result.contains("test_tool"));
assert!(result.contains("42"));
assert!(result.contains("custom_value"));
}
}