use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::time::Duration;
use tokio::sync::mpsc;
use crate::chat_job_queue::{ChatJobQueue, WebChatQueueDeps};
use crate::config::SharedAgentConfig;
use crate::cm_api_contract::chat::ConversationLayoutMeta;
use crate::conversation_store::{
self, CONVERSATION_STORE_MAX_ENTRIES, CONVERSATION_STORE_TTL_SECS, SaveConversationOutcome,
};
use crate::memory::long_term_memory::LongTermMemoryRuntime;
use crate::types::{CommandApprovalDecision, Message};
use crate::sse::SseStreamHub;
pub(crate) use crate::conversation_store::CONVERSATION_ID_MAX_LEN;
const CONVERSATION_STORE_TTL: Duration = Duration::from_secs(CONVERSATION_STORE_TTL_SECS);
pub(crate) struct ApprovalSessionSlot {
pub(crate) tx: mpsc::Sender<CommandApprovalDecision>,
pub(crate) created_at: std::time::Instant,
}
pub(crate) const APPROVAL_SESSION_TTL: Duration = Duration::from_secs(3600);
pub(crate) fn purge_expired_approval_sessions(
map: &mut HashMap<String, ApprovalSessionSlot>,
ttl: Duration,
) {
let now = std::time::Instant::now();
map.retain(|_, slot| now.duration_since(slot.created_at) <= ttl);
}
#[derive(Clone)]
pub(crate) struct MemoryConversationEntry {
messages: Vec<Message>,
active_agent_role: Option<String>,
active_session_mode: Option<String>,
layout: Option<ConversationLayoutMeta>,
revision: u64,
updated_at: std::time::Instant,
}
#[derive(Clone)]
pub(crate) struct ConversationTurnSeed {
pub messages: Vec<Message>,
pub expected_revision: Option<u64>,
pub persisted_active_agent_role: Option<String>,
pub persisted_active_session_mode: Option<String>,
pub layout: Option<ConversationLayoutMeta>,
}
fn nonempty_persisted_column(raw: String) -> Option<String> {
let t = raw.trim();
if t.is_empty() {
None
} else {
Some(t.to_string())
}
}
#[derive(Clone)]
pub(crate) struct AppStateHttpCore {
pub(crate) cfg: SharedAgentConfig,
pub(crate) config_path_for_reload: Option<String>,
pub(crate) api_key: Arc<str>,
pub(crate) client: reqwest::Client,
pub(crate) tools: Vec<crate::types::Tool>,
pub(crate) workspace_override: Arc<tokio::sync::RwLock<Option<String>>>,
pub(crate) uploads_dir: std::path::PathBuf,
}
pub(crate) async fn effective_workspace_path_from_override(
workspace_override: &tokio::sync::RwLock<Option<String>>,
cfg: &SharedAgentConfig,
) -> String {
let guard = workspace_override.read().await;
match guard.as_deref() {
None => String::new(),
Some(s) if s.trim().is_empty() => {
let cfg = cfg.read().await;
cfg.command_exec.run_command_working_dir.clone()
}
Some(s) => s.to_string(),
}
}
pub(crate) async fn workspace_is_set_from_override(
workspace_override: &tokio::sync::RwLock<Option<String>>,
) -> bool {
let guard = workspace_override.read().await;
guard.as_deref().is_some_and(|s| !s.trim().is_empty())
}
impl AppStateHttpCore {
pub(crate) async fn effective_workspace_path(&self) -> String {
effective_workspace_path_from_override(&self.workspace_override, &self.cfg).await
}
}
#[derive(Clone)]
pub(crate) struct AppStateChatRuntime {
pub(crate) chat_queue: ChatJobQueue,
pub(crate) chat_queue_job_deps: Arc<WebChatQueueDeps>,
}
#[derive(Clone)]
pub(crate) struct AppStateConversationRuntime {
pub(crate) conversation_backing: Arc<tokio::sync::RwLock<ConversationBacking>>,
pub(crate) conversation_id_counter: Arc<AtomicU64>,
}
#[derive(Clone)]
pub(crate) struct AppStateWebAux {
pub(crate) approval_sessions: Arc<tokio::sync::RwLock<HashMap<String, ApprovalSessionSlot>>>,
pub(crate) long_term_memory: Option<Arc<LongTermMemoryRuntime>>,
pub(crate) llm_models_health_cache:
Arc<std::sync::Mutex<Option<crate::health::CachedLlmModelsHealthProbe>>>,
pub(crate) sse_stream_hub: Arc<SseStreamHub>,
pub(crate) process_handles: Arc<crate::process_handles::ProcessHandles>,
pub(crate) async_chat_jobs: super::async_chat_job::AsyncChatJobsMap,
pub(crate) tool_job_registry: std::sync::Arc<crate::cm_internal::tool_jobs::ToolJobRegistry>,
pub(crate) mount_web_ui: bool,
}
#[derive(Clone)]
pub(crate) struct AppState {
pub(crate) http: AppStateHttpCore,
pub(crate) chat: AppStateChatRuntime,
pub(crate) conversation: AppStateConversationRuntime,
pub(crate) aux: AppStateWebAux,
}
#[derive(Clone)]
pub(crate) struct WebChatJobAppFacet {
pub(crate) conversation: AppStateConversationRuntime,
pub(crate) process_handles: Arc<crate::process_handles::TurnProcessHandles>,
pub(crate) approval_sessions: Arc<tokio::sync::RwLock<HashMap<String, ApprovalSessionSlot>>>,
pub(crate) tool_job_registry: std::sync::Arc<crate::cm_internal::tool_jobs::ToolJobRegistry>,
}
#[derive(Clone)]
pub(crate) enum ConversationBacking {
Memory(Arc<tokio::sync::RwLock<HashMap<String, MemoryConversationEntry>>>),
Sqlite(Arc<std::sync::Mutex<rusqlite::Connection>>),
}
impl ConversationBacking {
pub(crate) fn memory_default() -> Self {
Self::Memory(Arc::new(tokio::sync::RwLock::new(HashMap::new())))
}
pub(crate) fn is_sqlite(&self) -> bool {
matches!(self, Self::Sqlite(_))
}
}
async fn sqlite_conversation_store_op(
conn: Arc<std::sync::Mutex<rusqlite::Connection>>,
id_log: String,
op_zh: &'static str,
run: impl FnOnce(&rusqlite::Connection) -> Result<SaveConversationOutcome, rusqlite::Error>
+ Send
+ 'static,
) -> SaveConversationOutcome {
match tokio::task::spawn_blocking(move || {
let g = conn
.lock()
.map_err(|e: std::sync::PoisonError<_>| e.to_string())?;
run(&g).map_err(|e: rusqlite::Error| e.to_string())
})
.await
{
Ok(Ok(out)) => out,
Ok(Err(e)) => {
log::error!(
target: "crabmate",
"会话 SQLite {}失败 conversation_id={} error={}",
op_zh,
id_log,
e
);
SaveConversationOutcome::Conflict
}
Err(e) => {
log::error!(
target: "crabmate",
"会话 SQLite {}任务失败 conversation_id={} error={}",
op_zh,
id_log,
e
);
SaveConversationOutcome::Conflict
}
}
}
impl AppStateConversationRuntime {
pub(crate) async fn load_conversation_seed(
&self,
conversation_id: &str,
) -> Option<ConversationTurnSeed> {
let backing = self.conversation_backing.read().await;
match &*backing {
ConversationBacking::Memory(map) => {
let mut guard = map.write().await;
let entry = guard.get_mut(conversation_id)?;
if entry.updated_at.elapsed() > CONVERSATION_STORE_TTL {
guard.remove(conversation_id);
return None;
}
entry.updated_at = std::time::Instant::now();
Some(ConversationTurnSeed {
messages: entry.messages.clone(),
expected_revision: Some(entry.revision),
persisted_active_agent_role: entry.active_agent_role.clone(),
persisted_active_session_mode: entry.active_session_mode.clone(),
layout: entry.layout.clone(),
})
}
ConversationBacking::Sqlite(conn) => {
let id = conversation_id.to_string();
let c = Arc::clone(conn);
let loaded = tokio::task::spawn_blocking(move || {
let g = match c.lock() {
Ok(g) => g,
Err(e) => {
log::error!(
target: "crabmate",
"会话 SQLite 锁失败: {}",
e
);
return None;
}
};
match conversation_store::load(&g, &id, CONVERSATION_STORE_TTL_SECS) {
Ok(o) => o,
Err(e) => {
log::warn!(
target: "crabmate",
"会话 SQLite 读取失败 id={} error={}",
id,
e
);
None
}
}
})
.await
.ok()
.flatten();
loaded.map(|row| ConversationTurnSeed {
messages: row.messages,
expected_revision: Some(row.revision),
persisted_active_agent_role: nonempty_persisted_column(row.active_agent_role),
persisted_active_session_mode: nonempty_persisted_column(
row.active_session_mode,
),
layout: row.layout,
})
}
}
}
fn prune_memory_locked(
guard: &mut HashMap<String, MemoryConversationEntry>,
now: std::time::Instant,
) {
guard.retain(|_, v| now.duration_since(v.updated_at) <= CONVERSATION_STORE_TTL);
if guard.len() <= CONVERSATION_STORE_MAX_ENTRIES {
return;
}
let mut order: Vec<(String, std::time::Instant)> = guard
.iter()
.map(|(k, v)| (k.clone(), v.updated_at))
.collect();
order.sort_by_key(|(_, t)| *t);
let to_drop = guard.len() - CONVERSATION_STORE_MAX_ENTRIES;
for (k, _) in order.into_iter().take(to_drop) {
guard.remove(&k);
}
}
pub(crate) async fn save_conversation_messages_if_revision(
&self,
conversation_id: String,
messages: Vec<Message>,
active_agent_role: Option<&str>,
active_session_mode: Option<&str>,
expected_revision: Option<u64>,
) -> SaveConversationOutcome {
let backing = self.conversation_backing.read().await;
match &*backing {
ConversationBacking::Memory(map) => {
let mut guard = map.write().await;
let now = std::time::Instant::now();
let role_owned = active_agent_role
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_string);
let mode_owned = active_session_mode
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_string);
if let Some(entry) = guard.get_mut(&conversation_id) {
match expected_revision {
Some(exp) if entry.revision == exp => {
entry.layout = Some(crate::cm_turn_layout::layout_meta_from_messages(
&messages,
));
entry.messages = messages;
entry.active_agent_role = role_owned;
entry.active_session_mode = mode_owned;
entry.revision = entry.revision.saturating_add(1);
entry.updated_at = now;
}
_ => return SaveConversationOutcome::Conflict,
}
} else if expected_revision.is_some() {
return SaveConversationOutcome::Conflict;
} else {
let layout = Some(crate::cm_turn_layout::layout_meta_from_messages(&messages));
guard.insert(
conversation_id,
MemoryConversationEntry {
messages,
active_agent_role: role_owned,
active_session_mode: mode_owned,
layout,
revision: 1,
updated_at: now,
},
);
}
Self::prune_memory_locked(&mut guard, now);
SaveConversationOutcome::Saved
}
ConversationBacking::Sqlite(conn) => {
let id = conversation_id;
let id_log = id.clone();
let c = Arc::clone(conn);
let exp = expected_revision;
let active_for_sql = active_agent_role.map(|s| s.to_string());
let mode_for_sql = active_session_mode.map(|s| s.to_string());
sqlite_conversation_store_op(c, id_log, "保存", move |g| {
conversation_store::save_if_revision(
g,
&id,
messages,
active_for_sql.as_deref(),
mode_for_sql.as_deref(),
exp,
)
})
.await
}
}
}
pub(crate) async fn truncate_conversation_before_user_ordinal_if_revision(
&self,
conversation_id: String,
user_ordinal: usize,
expected_revision: u64,
) -> SaveConversationOutcome {
let backing = self.conversation_backing.read().await;
match &*backing {
ConversationBacking::Memory(map) => {
let mut guard = map.write().await;
let Some(entry) = guard.get_mut(&conversation_id) else {
return SaveConversationOutcome::Conflict;
};
if entry.updated_at.elapsed() > CONVERSATION_STORE_TTL {
guard.remove(&conversation_id);
return SaveConversationOutcome::Conflict;
}
if entry.revision != expected_revision {
return SaveConversationOutcome::Conflict;
}
let mut u = 0usize;
let mut cut = entry.messages.len();
for (i, m) in entry.messages.iter().enumerate() {
if crate::types::user_message_counts_for_branch_truncation(m) {
if u == user_ordinal {
cut = i;
break;
}
u += 1;
}
}
if cut >= entry.messages.len() {
entry.updated_at = std::time::Instant::now();
return SaveConversationOutcome::Saved;
}
entry.messages.truncate(cut);
entry.layout = Some(crate::cm_turn_layout::layout_meta_from_messages(
&entry.messages,
));
entry.revision = entry.revision.saturating_add(1);
entry.updated_at = std::time::Instant::now();
Self::prune_memory_locked(&mut guard, std::time::Instant::now());
SaveConversationOutcome::Saved
}
ConversationBacking::Sqlite(conn) => {
let id = conversation_id;
let id_log = id.clone();
let c = Arc::clone(conn);
sqlite_conversation_store_op(c, id_log, "截断", move |g| {
conversation_store::truncate_before_user_ordinal_if_revision(
g,
&id,
user_ordinal,
expected_revision,
)
})
.await
}
}
}
pub(crate) async fn referenced_upload_filenames(&self) -> HashSet<String> {
use crate::web::chat_uploads_paths::collect_upload_filenames_from_text;
let backing = self.conversation_backing.read().await;
match &*backing {
ConversationBacking::Memory(map) => {
let guard = map.read().await;
let mut out = HashSet::new();
for entry in guard.values() {
if let Ok(s) = serde_json::to_string(&entry.messages) {
out.extend(collect_upload_filenames_from_text(&s));
}
}
out
}
ConversationBacking::Sqlite(conn) => {
let c = Arc::clone(conn);
tokio::task::spawn_blocking(move || {
let g = match c.lock() {
Ok(g) => g,
Err(_) => return HashSet::new(),
};
let Ok(blobs) = conversation_store::list_all_messages_json(&g) else {
return HashSet::new();
};
let mut out = HashSet::new();
for s in blobs {
out.extend(collect_upload_filenames_from_text(&s));
}
out
})
.await
.unwrap_or_default()
}
}
}
pub(crate) async fn conversation_count(&self) -> usize {
let backing = self.conversation_backing.read().await;
match &*backing {
ConversationBacking::Memory(map) => map.read().await.len(),
ConversationBacking::Sqlite(conn) => {
let c = Arc::clone(conn);
tokio::task::spawn_blocking(move || {
let g = match c.lock() {
Ok(g) => g,
Err(_) => return 0usize,
};
conversation_store::count(&g).unwrap_or(0)
})
.await
.unwrap_or(0)
}
}
}
pub(crate) async fn delete_conversation_record(&self, conversation_id: &str) {
let backing = self.conversation_backing.read().await;
match &*backing {
ConversationBacking::Memory(map) => {
let mut guard = map.write().await;
guard.remove(conversation_id);
}
ConversationBacking::Sqlite(conn) => {
let id = conversation_id.to_string();
let c = Arc::clone(conn);
let _ = tokio::task::spawn_blocking(move || {
if let Ok(g) = c.lock() {
let _ = conversation_store::delete_by_id(&g, &id);
}
})
.await;
}
}
}
}
pub(crate) fn open_conversation_sqlite(
path: &Path,
) -> Result<Arc<std::sync::Mutex<rusqlite::Connection>>, Box<dyn std::error::Error + Send + Sync>> {
let conn = conversation_store::open_file(path)?;
if let Err(e) = LongTermMemoryRuntime::migrate_on_connection(&conn) {
return Err(format!("长期记忆表迁移失败: {e}").into());
}
Ok(Arc::new(std::sync::Mutex::new(conn)))
}