434 lines
14 KiB
Rust
434 lines
14 KiB
Rust
use crate::agent::{SystemPrompt, SystemPromptContext, SystemPromptProvider};
|
|
use crate::command::context::CommandContext;
|
|
use crate::command::handler::CommandHandler;
|
|
use crate::command::response::{CommandError, CommandResponse, MessageKind};
|
|
use crate::command::Command;
|
|
use crate::config::LLMProviderConfig;
|
|
use crate::gateway::agent_prompt_provider::SimpleAgentPromptProvider;
|
|
use crate::storage::{SessionRecord, SessionStore};
|
|
use async_trait::async_trait;
|
|
use chrono::{Local, TimeZone};
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
|
|
/// 保存会话命令处理器
|
|
///
|
|
/// 将当前会话内容（系统提示词和消息历史）保存到 Markdown 文件
|
|
pub struct SaveSessionCommandHandler {
|
|
store: Arc<SessionStore>,
|
|
provider_config: LLMProviderConfig,
|
|
}
|
|
|
|
impl SaveSessionCommandHandler {
|
|
/// 创建新的保存会话命令处理器
|
|
///
|
|
/// # Arguments
|
|
/// * `store` - 会话存储
|
|
/// * `provider_config` - LLM 提供者配置（用于构建系统提示词）
|
|
pub fn new(store: Arc<SessionStore>, provider_config: LLMProviderConfig) -> Self {
|
|
Self {
|
|
store,
|
|
provider_config,
|
|
}
|
|
}
|
|
|
|
/// 从会话记录获取存储（用于测试）
|
|
#[cfg(test)]
|
|
fn store(&self) -> &Arc<SessionStore> {
|
|
&self.store
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl CommandHandler for SaveSessionCommandHandler {
|
|
fn can_handle(&self, cmd: &Command) -> bool {
|
|
matches!(cmd, Command::SaveSession { .. })
|
|
}
|
|
|
|
async fn handle(
|
|
&self,
|
|
cmd: Command,
|
|
ctx: CommandContext,
|
|
) -> Result<CommandResponse, CommandError> {
|
|
match cmd {
|
|
Command::SaveSession { filepath } => {
|
|
handle_save_session(self, filepath, ctx).await
|
|
}
|
|
_ => unreachable!(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 处理保存会话命令
|
|
async fn handle_save_session(
|
|
handler: &SaveSessionCommandHandler,
|
|
filepath: Option<String>,
|
|
ctx: CommandContext,
|
|
) -> Result<CommandResponse, CommandError> {
|
|
let session_id = ctx
|
|
.session_id
|
|
.as_deref()
|
|
.ok_or_else(|| CommandError::new("NO_SESSION", "No active session".to_string()))?;
|
|
|
|
// 获取会话记录
|
|
let record = handler
|
|
.store
|
|
.get_session(session_id)
|
|
.map_err(|e| CommandError::new("SESSION_ERROR", e.to_string()))?
|
|
.ok_or_else(|| CommandError::new("SESSION_NOT_FOUND", "Session not found".to_string()))?;
|
|
|
|
// 获取所有消息（包括历史）
|
|
let messages = handler
|
|
.store
|
|
.load_all_messages(session_id)
|
|
.map_err(|e| CommandError::new("LOAD_MESSAGES_ERROR", e.to_string()))?;
|
|
|
|
// 计算用户消息数（用于系统提示词构建）
|
|
let user_message_count = messages.iter().filter(|m| m.role == "user").count();
|
|
|
|
// 构建系统提示词
|
|
let system_prompt = build_system_prompt(&handler.provider_config, &record, user_message_count);
|
|
|
|
// 生成 Markdown 内容
|
|
let markdown = generate_markdown(&record, &system_prompt, &messages);
|
|
|
|
// 确定输出路径
|
|
let output_path = resolve_filepath(filepath, &record);
|
|
|
|
// 创建父目录
|
|
if let Some(parent) = output_path.parent() {
|
|
if !parent.as_os_str().is_empty() && !parent.exists() {
|
|
std::fs::create_dir_all(parent).map_err(|e| {
|
|
CommandError::new(
|
|
"CREATE_DIR_ERROR",
|
|
format!("Failed to create directory: {}", e),
|
|
)
|
|
})?;
|
|
}
|
|
}
|
|
|
|
// 写入文件
|
|
std::fs::write(&output_path, markdown).map_err(|e| {
|
|
CommandError::new(
|
|
"WRITE_FILE_ERROR",
|
|
format!("Failed to write file: {}", e),
|
|
)
|
|
})?;
|
|
|
|
Ok(CommandResponse::success(ctx.request_id)
|
|
.with_message(
|
|
MessageKind::Notification,
|
|
&format!("Session saved to: {}", output_path.display()),
|
|
)
|
|
.with_metadata("filepath", output_path.to_string_lossy().as_ref())
|
|
.with_metadata("message_count", &messages.len().to_string()))
|
|
}
|
|
|
|
/// 构建系统提示词
|
|
fn build_system_prompt(
|
|
provider_config: &LLMProviderConfig,
|
|
record: &SessionRecord,
|
|
user_message_count: usize,
|
|
) -> Option<SystemPrompt> {
|
|
let provider = SimpleAgentPromptProvider::new(provider_config.clone());
|
|
let context = SystemPromptContext {
|
|
session_id: Some(record.id.clone()),
|
|
chat_id: record.chat_id.clone(),
|
|
user_message_count,
|
|
};
|
|
provider.build(&context)
|
|
}
|
|
|
|
/// 生成 Markdown 内容
|
|
fn generate_markdown(
|
|
record: &SessionRecord,
|
|
system_prompt: &Option<SystemPrompt>,
|
|
messages: &[crate::bus::ChatMessage],
|
|
) -> String {
|
|
let mut output = String::new();
|
|
|
|
// YAML frontmatter
|
|
output.push_str("---\n");
|
|
output.push_str(&format!("title: {}\n", escape_yaml_string(&record.title)));
|
|
output.push_str(&format!("session_id: {}\n", record.id));
|
|
output.push_str(&format!("channel: {}\n", record.channel_name));
|
|
output.push_str(&format!("chat_id: {}\n", record.chat_id));
|
|
output.push_str(&format!(
|
|
"created_at: {}\n",
|
|
format_timestamp(record.created_at)
|
|
));
|
|
output.push_str(&format!(
|
|
"updated_at: {}\n",
|
|
format_timestamp(record.updated_at)
|
|
));
|
|
output.push_str(&format!(
|
|
"last_active_at: {}\n",
|
|
format_timestamp(record.last_active_at)
|
|
));
|
|
output.push_str(&format!("message_count: {}\n", messages.len()));
|
|
output.push_str("---\n\n");
|
|
|
|
// 系统提示词
|
|
output.push_str("# System Prompt\n\n");
|
|
if let Some(prompt) = system_prompt {
|
|
output.push_str("```\n");
|
|
output.push_str(&prompt.content);
|
|
output.push_str("\n```\n\n");
|
|
} else {
|
|
output.push_str("*No system prompt available*\n\n");
|
|
}
|
|
|
|
// 消息历史
|
|
output.push_str("# Message History\n\n");
|
|
|
|
for (idx, msg) in messages.iter().enumerate() {
|
|
output.push_str(&format!("## Message {}\n\n", idx + 1));
|
|
output.push_str(&format!("**Role:** {}\n\n", msg.role));
|
|
output.push_str(&format!("**ID:** {}\n\n", msg.id));
|
|
output.push_str(&format!(
|
|
"**Time:** {}\n\n",
|
|
format_timestamp(msg.timestamp)
|
|
));
|
|
|
|
if let Some(ref ctx) = msg.system_context {
|
|
output.push_str(&format!("**System Context:** `{}`\n\n", ctx));
|
|
}
|
|
|
|
if let Some(ref tool_name) = msg.tool_name {
|
|
output.push_str(&format!("**Tool Name:** `{}`\n\n", tool_name));
|
|
}
|
|
|
|
if let Some(ref tool_call_id) = msg.tool_call_id {
|
|
output.push_str(&format!("**Tool Call ID:** `{}`\n\n", tool_call_id));
|
|
}
|
|
|
|
if let Some(ref reasoning) = msg.reasoning_content {
|
|
output.push_str("### Reasoning\n\n");
|
|
output.push_str("```\n");
|
|
output.push_str(reasoning);
|
|
output.push_str("\n```\n\n");
|
|
}
|
|
|
|
// Content
|
|
output.push_str("### Content\n\n");
|
|
if msg.content.is_empty() {
|
|
output.push_str("*empty*\n\n");
|
|
} else {
|
|
output.push_str(&format!("{}\n\n", format_message_content(&msg.content)));
|
|
}
|
|
|
|
// Tool calls
|
|
if let Some(ref calls) = msg.tool_calls {
|
|
if !calls.is_empty() {
|
|
output.push_str("### Tool Calls\n\n");
|
|
for call in calls {
|
|
output.push_str(&format!("- **{}** (`{}`)\n", call.name, call.id));
|
|
output.push_str(" ```json\n");
|
|
let args_json = serde_json::to_string_pretty(&call.arguments)
|
|
.unwrap_or_else(|_| call.arguments.to_string());
|
|
for line in args_json.lines() {
|
|
output.push_str(&format!(" {}\n", line));
|
|
}
|
|
output.push_str(" ```\n");
|
|
}
|
|
output.push('\n');
|
|
}
|
|
}
|
|
|
|
// Media refs
|
|
if !msg.media_refs.is_empty() {
|
|
output.push_str("### Media References\n\n");
|
|
for media_ref in &msg.media_refs {
|
|
output.push_str(&format!("- `{}`\n", media_ref));
|
|
}
|
|
output.push('\n');
|
|
}
|
|
|
|
output.push_str("---\n\n");
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// 格式化消息内容
|
|
///
|
|
/// 如果内容包含特殊字符，使用代码块包装
|
|
fn format_message_content(content: &str) -> String {
|
|
// 如果内容包含代码块标记或表格标记，使用原始格式
|
|
if content.contains("```") || content.contains("| ") {
|
|
format!("```\n{}\n```", content)
|
|
} else {
|
|
content.to_string()
|
|
}
|
|
}
|
|
|
|
/// 转义 YAML 字符串
|
|
fn escape_yaml_string(s: &str) -> String {
|
|
if s.contains('\n') || s.contains('"') || s.contains(':') || s.starts_with(' ') {
|
|
// 使用双引号包裹并转义内部的双引号
|
|
format!("\"{}\"", s.replace('"', "\\\""))
|
|
} else {
|
|
s.to_string()
|
|
}
|
|
}
|
|
|
|
/// 格式化时间戳
|
|
fn format_timestamp(ts: i64) -> String {
|
|
Local
|
|
.timestamp_millis_opt(ts)
|
|
.single()
|
|
.map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
|
|
.unwrap_or_else(|| format!("{}", ts))
|
|
}
|
|
|
|
/// 解析文件路径
|
|
///
|
|
/// 如果未提供路径，自动生成基于会话标题和时间戳的文件名
|
|
fn resolve_filepath(filepath: Option<String>, record: &SessionRecord) -> PathBuf {
|
|
match filepath {
|
|
Some(path) => PathBuf::from(path),
|
|
None => {
|
|
// 生成安全标题（替换特殊字符）
|
|
let safe_title = record
|
|
.title
|
|
.replace(' ', "_")
|
|
.replace('/', "_")
|
|
.replace('\\', "_")
|
|
.replace(':', "_")
|
|
.replace('<', "_")
|
|
.replace('>', "_")
|
|
.replace('|', "_")
|
|
.replace('?', "_")
|
|
.replace('*', "_")
|
|
.replace('"', "_");
|
|
|
|
// 使用标题或 session_id 作为文件名
|
|
let base_name = if safe_title.is_empty() {
|
|
format!("session_{}", &record.id[..8.min(record.id.len())])
|
|
} else {
|
|
safe_title
|
|
};
|
|
|
|
// 添加时间戳
|
|
let timestamp = Local::now().format("%Y%m%d_%H%M%S");
|
|
let filename = format!("{}_{}.md", base_name, timestamp);
|
|
|
|
PathBuf::from(filename)
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::storage::{SessionRecord, SessionStore};
|
|
use std::collections::HashMap;
|
|
|
|
fn test_config() -> LLMProviderConfig {
|
|
LLMProviderConfig {
|
|
provider_type: "openai".to_string(),
|
|
name: "test".to_string(),
|
|
base_url: "http://localhost".to_string(),
|
|
api_key: "test-key".to_string(),
|
|
extra_headers: HashMap::new(),
|
|
llm_timeout_secs: 120,
|
|
memory_maintenance_timeout_secs: 600,
|
|
model_id: "test-model".to_string(),
|
|
temperature: Some(0.0),
|
|
max_tokens: Some(32),
|
|
context_window_tokens: None,
|
|
tool_result_max_chars: 20_000,
|
|
context_tool_result_trim_chars: 20_000,
|
|
model_extra: HashMap::new(),
|
|
max_tool_iterations: 1,
|
|
}
|
|
}
|
|
|
|
fn create_test_record(id: &str, title: &str) -> SessionRecord {
|
|
SessionRecord {
|
|
id: id.to_string(),
|
|
title: title.to_string(),
|
|
channel_name: "cli".to_string(),
|
|
chat_id: id.to_string(),
|
|
summary: None,
|
|
created_at: 1705312800000, // 2024-01-15 10:00:00
|
|
updated_at: 1705316400000, // 2024-01-15 11:00:00
|
|
last_active_at: 1705316400000,
|
|
archived_at: None,
|
|
deleted_at: None,
|
|
message_count: 0,
|
|
reset_cutoff_seq: 0,
|
|
user_turn_count: 0,
|
|
agent_prompt_reinjection_count: 0,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_resolve_filepath_with_custom_path() {
|
|
let record = create_test_record("test-123", "My Session");
|
|
let path = resolve_filepath(Some("/custom/path/file.md".to_string()), &record);
|
|
assert_eq!(path, PathBuf::from("/custom/path/file.md"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_resolve_filepath_generates_filename_with_title() {
|
|
let record = create_test_record("test-123", "My Session");
|
|
let path = resolve_filepath(None, &record);
|
|
let filename = path.file_name().unwrap().to_str().unwrap();
|
|
|
|
assert!(filename.starts_with("My_Session_"));
|
|
assert!(filename.ends_with(".md"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_resolve_filepath_generates_filename_with_id_when_title_empty() {
|
|
let record = create_test_record("abc12345-xyz", "");
|
|
let path = resolve_filepath(None, &record);
|
|
let filename = path.file_name().unwrap().to_str().unwrap();
|
|
|
|
assert!(filename.starts_with("session_abc123"));
|
|
assert!(filename.ends_with(".md"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_escape_yaml_string() {
|
|
assert_eq!(escape_yaml_string("simple"), "simple");
|
|
assert_eq!(escape_yaml_string("with: colon"), "\"with: colon\"");
|
|
assert_eq!(escape_yaml_string("with \"quote\""), "\"with \\\"quote\\\"\"");
|
|
}
|
|
|
|
#[test]
|
|
fn test_format_message_content() {
|
|
assert_eq!(format_message_content("hello"), "hello");
|
|
assert_eq!(
|
|
format_message_content("```code```"),
|
|
"```\n```code```\n```"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_markdown_structure() {
|
|
let record = create_test_record("test-123", "Test Session");
|
|
let messages = vec![crate::bus::ChatMessage::system("System prompt here")];
|
|
|
|
let markdown = generate_markdown(&record, &None, &messages);
|
|
|
|
assert!(markdown.contains("---"));
|
|
assert!(markdown.contains("title:"));
|
|
assert!(markdown.contains("session_id: test-123"));
|
|
assert!(markdown.contains("# System Prompt"));
|
|
assert!(markdown.contains("# Message History"));
|
|
assert!(markdown.contains("## Message 1"));
|
|
assert!(markdown.contains("**Role:** system"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_can_handle() {
|
|
let store = Arc::new(SessionStore::in_memory().unwrap());
|
|
let handler = SaveSessionCommandHandler::new(store, test_config());
|
|
|
|
assert!(handler.can_handle(&Command::SaveSession { filepath: None }));
|
|
assert!(!handler.can_handle(&Command::CreateSession { title: None }));
|
|
}
|
|
}
|