refactor: split main.rs into 7 modules, add life loop with timer system
Structure: main.rs (534) — entry, handler, prompt building config.rs (52) — config structs state.rs (358) — AppState, SQLite, persistence tools.rs (665) — tool definitions, execution, subagent management stream.rs (776) — OpenAI/Claude streaming, system prompt display.rs (220)— markdown rendering, message formatting life.rs (87) — life loop heartbeat, timer firing New features: - Life Loop: background tokio task, 30s heartbeat, scans timers table - Timer tools: set_timer (relative/absolute/cron), list_timers, cancel_timer - inner_state table for life loop's own context - cron crate for recurring schedule parsing Zero logic changes in the refactor — pure structural split.
This commit is contained in:
Fam Zheng
358
src/state.rs
Normal file
358
src/state.rs
Normal file
@@ -0,0 +1,358 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::Result;
|
||||
use chrono::NaiveDate;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{error, info};
|
||||
|
||||
use crate::tools::SubAgent;
|
||||
|
||||
// ── persistent state ────────────────────────────────────────────────
|
||||
|
||||
#[derive(Serialize, Deserialize, Default)]
|
||||
pub struct Persistent {
|
||||
pub authed: HashMap<i64, NaiveDate>,
|
||||
pub known_sessions: HashSet<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Clone, Default)]
|
||||
pub struct ConversationState {
|
||||
pub summary: String,
|
||||
pub messages: Vec<serde_json::Value>,
|
||||
pub total_messages: usize,
|
||||
}
|
||||
|
||||
pub const MAX_WINDOW: usize = 100;
|
||||
pub const SLIDE_SIZE: usize = 50;
|
||||
|
||||
pub struct AppState {
|
||||
pub persist: RwLock<Persistent>,
|
||||
pub state_path: PathBuf,
|
||||
pub db: tokio::sync::Mutex<rusqlite::Connection>,
|
||||
pub agents: RwLock<HashMap<String, Arc<SubAgent>>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn load(path: PathBuf) -> Self {
|
||||
let persist = std::fs::read_to_string(&path)
|
||||
.ok()
|
||||
.and_then(|s| serde_json::from_str(&s).ok())
|
||||
.unwrap_or_default();
|
||||
info!("loaded state from {}", path.display());
|
||||
|
||||
let db_path = path.parent().unwrap_or(Path::new(".")).join("noc.db");
|
||||
let conn = rusqlite::Connection::open(&db_path)
|
||||
.unwrap_or_else(|e| panic!("open {}: {e}", db_path.display()));
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE IF NOT EXISTS conversations (
|
||||
session_id TEXT PRIMARY KEY,
|
||||
summary TEXT NOT NULL DEFAULT '',
|
||||
total_messages INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
session_id TEXT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now', 'localtime'))
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);
|
||||
CREATE TABLE IF NOT EXISTS scratch_area (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
content TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS config (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL DEFAULT '',
|
||||
create_time TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
update_time TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS config_history (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
create_time TEXT NOT NULL,
|
||||
update_time TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS memory_slots (
|
||||
slot_nr INTEGER PRIMARY KEY CHECK(slot_nr BETWEEN 0 AND 99),
|
||||
content TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS timers (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
chat_id INTEGER NOT NULL,
|
||||
label TEXT NOT NULL,
|
||||
schedule TEXT NOT NULL,
|
||||
next_fire TEXT NOT NULL,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now', 'localtime'))
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS inner_state (
|
||||
id INTEGER PRIMARY KEY CHECK(id = 1),
|
||||
content TEXT NOT NULL DEFAULT ''
|
||||
);
|
||||
INSERT OR IGNORE INTO inner_state (id, content) VALUES (1, '');",
|
||||
)
|
||||
.expect("init db schema");
|
||||
|
||||
// migrations
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE messages ADD COLUMN created_at TEXT NOT NULL DEFAULT ''",
|
||||
[],
|
||||
);
|
||||
|
||||
info!("opened db {}", db_path.display());
|
||||
|
||||
Self {
|
||||
persist: RwLock::new(persist),
|
||||
state_path: path,
|
||||
db: tokio::sync::Mutex::new(conn),
|
||||
agents: RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn save(&self) {
|
||||
let data = self.persist.read().await;
|
||||
if let Ok(json) = serde_json::to_string_pretty(&*data) {
|
||||
if let Err(e) = std::fs::write(&self.state_path, json) {
|
||||
error!("save state: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn load_conv(&self, sid: &str) -> ConversationState {
|
||||
let db = self.db.lock().await;
|
||||
let (summary, total) = db
|
||||
.query_row(
|
||||
"SELECT summary, total_messages FROM conversations WHERE session_id = ?1",
|
||||
[sid],
|
||||
|row| Ok((row.get::<_, String>(0)?, row.get::<_, usize>(1)?)),
|
||||
)
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut stmt = db
|
||||
.prepare("SELECT role, content, created_at FROM messages WHERE session_id = ?1 ORDER BY id")
|
||||
.unwrap();
|
||||
let messages: Vec<serde_json::Value> = stmt
|
||||
.query_map([sid], |row| {
|
||||
let role: String = row.get(0)?;
|
||||
let content: String = row.get(1)?;
|
||||
let ts: String = row.get(2)?;
|
||||
let tagged = if ts.is_empty() {
|
||||
content
|
||||
} else {
|
||||
format!("[{ts}] {content}")
|
||||
};
|
||||
Ok(serde_json::json!({"role": role, "content": tagged}))
|
||||
})
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
|
||||
ConversationState {
|
||||
summary,
|
||||
messages,
|
||||
total_messages: total,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn push_message(&self, sid: &str, role: &str, content: &str) {
|
||||
let db = self.db.lock().await;
|
||||
let _ = db.execute(
|
||||
"INSERT OR IGNORE INTO conversations (session_id) VALUES (?1)",
|
||||
[sid],
|
||||
);
|
||||
let _ = db.execute(
|
||||
"INSERT INTO messages (session_id, role, content, created_at) VALUES (?1, ?2, ?3, datetime('now', 'localtime'))",
|
||||
rusqlite::params![sid, role, content],
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn message_count(&self, sid: &str) -> usize {
|
||||
let db = self.db.lock().await;
|
||||
db.query_row(
|
||||
"SELECT COUNT(*) FROM messages WHERE session_id = ?1",
|
||||
[sid],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub async fn slide_window(&self, sid: &str, new_summary: &str, slide_size: usize) {
|
||||
let db = self.db.lock().await;
|
||||
let _ = db.execute(
|
||||
"DELETE FROM messages WHERE id IN (
|
||||
SELECT id FROM messages WHERE session_id = ?1 ORDER BY id LIMIT ?2
|
||||
)",
|
||||
rusqlite::params![sid, slide_size],
|
||||
);
|
||||
let _ = db.execute(
|
||||
"UPDATE conversations SET summary = ?1, total_messages = total_messages + ?2 \
|
||||
WHERE session_id = ?3",
|
||||
rusqlite::params![new_summary, slide_size, sid],
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn get_oldest_messages(&self, sid: &str, count: usize) -> Vec<serde_json::Value> {
|
||||
let db = self.db.lock().await;
|
||||
let mut stmt = db
|
||||
.prepare(
|
||||
"SELECT role, content FROM messages WHERE session_id = ?1 ORDER BY id LIMIT ?2",
|
||||
)
|
||||
.unwrap();
|
||||
stmt.query_map(rusqlite::params![sid, count], |row| {
|
||||
let role: String = row.get(0)?;
|
||||
let content: String = row.get(1)?;
|
||||
Ok(serde_json::json!({"role": role, "content": content}))
|
||||
})
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn get_scratch(&self) -> String {
|
||||
let db = self.db.lock().await;
|
||||
db.query_row(
|
||||
"SELECT content FROM scratch_area ORDER BY id DESC LIMIT 1",
|
||||
[],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub async fn push_scratch(&self, content: &str) {
|
||||
let db = self.db.lock().await;
|
||||
let _ = db.execute(
|
||||
"INSERT INTO scratch_area (content) VALUES (?1)",
|
||||
[content],
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn get_config(&self, key: &str) -> Option<String> {
|
||||
let db = self.db.lock().await;
|
||||
db.query_row(
|
||||
"SELECT value FROM config WHERE key = ?1",
|
||||
[key],
|
||||
|row| row.get(0),
|
||||
)
|
||||
.ok()
|
||||
}
|
||||
|
||||
pub async fn get_inner_state(&self) -> String {
|
||||
let db = self.db.lock().await;
|
||||
db.query_row("SELECT content FROM inner_state WHERE id = 1", [], |row| row.get(0))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[allow(dead_code)] // used by life loop tools (coming soon)
|
||||
pub async fn set_inner_state(&self, content: &str) {
|
||||
let db = self.db.lock().await;
|
||||
let _ = db.execute(
|
||||
"UPDATE inner_state SET content = ?1 WHERE id = 1",
|
||||
[content],
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn add_timer(&self, chat_id: i64, label: &str, schedule: &str, next_fire: &str) -> i64 {
|
||||
let db = self.db.lock().await;
|
||||
db.execute(
|
||||
"INSERT INTO timers (chat_id, label, schedule, next_fire) VALUES (?1, ?2, ?3, ?4)",
|
||||
rusqlite::params![chat_id, label, schedule, next_fire],
|
||||
)
|
||||
.unwrap();
|
||||
db.last_insert_rowid()
|
||||
}
|
||||
|
||||
pub async fn list_timers(&self, chat_id: Option<i64>) -> Vec<(i64, i64, String, String, String, bool)> {
|
||||
let db = self.db.lock().await;
|
||||
let (sql, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match chat_id {
|
||||
Some(cid) => (
|
||||
"SELECT id, chat_id, label, schedule, next_fire, enabled FROM timers WHERE chat_id = ?1 ORDER BY next_fire",
|
||||
vec![Box::new(cid)],
|
||||
),
|
||||
None => (
|
||||
"SELECT id, chat_id, label, schedule, next_fire, enabled FROM timers ORDER BY next_fire",
|
||||
vec![],
|
||||
),
|
||||
};
|
||||
let mut stmt = db.prepare(sql).unwrap();
|
||||
stmt.query_map(rusqlite::params_from_iter(params), |row| {
|
||||
Ok((
|
||||
row.get(0)?,
|
||||
row.get(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
row.get::<_, String>(4)?,
|
||||
row.get::<_, bool>(5)?,
|
||||
))
|
||||
})
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn cancel_timer(&self, timer_id: i64) -> bool {
|
||||
let db = self.db.lock().await;
|
||||
db.execute("DELETE FROM timers WHERE id = ?1", [timer_id]).unwrap() > 0
|
||||
}
|
||||
|
||||
pub async fn due_timers(&self) -> Vec<(i64, i64, String, String)> {
|
||||
let db = self.db.lock().await;
|
||||
let now = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let mut stmt = db
|
||||
.prepare(
|
||||
"SELECT id, chat_id, label, schedule FROM timers WHERE enabled = 1 AND next_fire <= ?1",
|
||||
)
|
||||
.unwrap();
|
||||
stmt.query_map([&now], |row| {
|
||||
Ok((
|
||||
row.get(0)?,
|
||||
row.get(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
})
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn update_timer_next_fire(&self, timer_id: i64, next_fire: &str) {
|
||||
let db = self.db.lock().await;
|
||||
let _ = db.execute(
|
||||
"UPDATE timers SET next_fire = ?1 WHERE id = ?2",
|
||||
rusqlite::params![next_fire, timer_id],
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn get_memory_slots(&self) -> Vec<(i32, String)> {
|
||||
let db = self.db.lock().await;
|
||||
let mut stmt = db
|
||||
.prepare("SELECT slot_nr, content FROM memory_slots WHERE content != '' ORDER BY slot_nr")
|
||||
.unwrap();
|
||||
stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn set_memory_slot(&self, slot_nr: i32, content: &str) -> Result<()> {
|
||||
if !(0..=99).contains(&slot_nr) {
|
||||
anyhow::bail!("slot_nr must be 0-99, got {slot_nr}");
|
||||
}
|
||||
if content.len() > 200 {
|
||||
anyhow::bail!("content too long: {} chars (max 200)", content.len());
|
||||
}
|
||||
let db = self.db.lock().await;
|
||||
db.execute(
|
||||
"INSERT INTO memory_slots (slot_nr, content) VALUES (?1, ?2) \
|
||||
ON CONFLICT(slot_nr) DO UPDATE SET content = ?2",
|
||||
rusqlite::params![slot_nr, content],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user