iAs/src/core/session.rs

use crate::llm::types::Message;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;

/// 摘要原因
#[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 user_id: String,
    pub db_pool: Option<Arc<sqlx::PgPool>>,
    /// 当前正在处理消息的用户 ID（spawn 前设置，executor 中读取）
    pub current_user_id: String,
    /// 摘要总结点——此位置之前的消息已被压缩
    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 {
            user_id: "default".to_string(),
            db_pool: None,
            current_user_id: String::new(),
            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 latest_overflow_summary(&self) -> Option<&SummaryEntry> {
        self.summaries
            .iter()
            .rev()
            .find(|s| s.reason == SummaryReason::Overflow)
    }

    /// 从预载的历史记录初始化会话（Worker 模式）
    pub fn load_from_history(&mut self, history: &[crate::ipc::HistoryEntry]) {
        for entry in history {
            match entry.role.as_str() {
                "user" => self.add_user(entry.content.clone()),
                "assistant" => self.add_assistant(entry.content.clone()),
                _ => {}
            }
        }
        if !history.is_empty() {
            self.last_user_at = Some(chrono::Utc::now());
        }
    }

    /// 从数据库加载最近的聊天记录到会话中
    pub async fn load_recent_messages(&mut self, limit: i64) {
        let pool = match &self.db_pool {
            Some(p) => p.clone(),
            None => return,
        };

        let uid = self.current_user_id.clone();
        let rows = if uid.is_empty() {
            sqlx::query_as::<_, (String, String, String, String)>(
                "SELECT direction, user_id, text, message_id FROM chat_records ORDER BY created_at DESC LIMIT $1",
            )
            .bind(limit)
            .fetch_all(pool.as_ref())
            .await
        } else {
            sqlx::query_as::<_, (String, String, String, String)>(
                "SELECT direction, user_id, text, message_id FROM chat_records WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2",
            )
            .bind(&uid)
            .bind(limit)
            .fetch_all(pool.as_ref())
            .await
        };

        if let Ok(rows) = rows {
            let msgs: Vec<Message> = rows
                .into_iter()
                .rev() // 恢复时间顺序
                .map(|(dir, _uid, text, _mid)| {
                    if dir == "inbound" {
                        Message::user(text)
                    } else {
                        Message::assistant(text)
                    }
                })
                .collect();
            if !msgs.is_empty() {
                tracing::info!("从数据库恢复 {} 条历史消息", msgs.len());
                self.messages = msgs;
                self.checkpoint = 0;
                self.last_user_at = Some(Utc::now());
            }
        }
    }

    /// 保存摘要到数据库（委托给 db::models::save_summary）
    pub async fn save_summary_to_db(&self, text: &str, reason: &str, message_count: i32) {
        let Some(pool) = &self.db_pool else { return };
        let uid = if self.current_user_id.is_empty() {
            if self.user_id.is_empty() { "default" } else { &self.user_id }
        } else {
            &self.current_user_id
        };
        if let Err(e) = crate::db::models::save_summary(pool, uid, text, reason, message_count).await {
            tracing::warn!("保存摘要到数据库失败: {}", e);
        }
    }
}