damocles-daemon/src/main.rs

801 lines
25 KiB
Rust

mod paths;
use std::fmt::Write as _;
use std::path::Path;
use std::sync::Arc;
use anyhow::{Context, bail};
use matrix_sdk::{
Client, Room, RoomState,
authentication::matrix::MatrixSession,
config::SyncSettings,
ruma::{
OwnedEventId, OwnedRoomId, OwnedUserId, UserId,
api::client::{
filter::FilterDefinition, receipt::create_receipt::v3::ReceiptType as CreateReceiptType,
},
events::{
receipt::ReceiptThread,
room::{
member::StrippedRoomMemberEvent,
message::{MessageType, OriginalSyncRoomMessageEvent, RoomMessageEventContent},
},
},
},
};
use serde::{Deserialize, Serialize};
use tokio::fs;
use tokio::sync::Mutex;
use tracing_subscriber::EnvFilter;
#[derive(Debug, Deserialize)]
struct Config {
homeserver: String,
username: String,
password: String,
rate_limit_per_min: Option<u32>,
model: Option<String>,
}
const DEFAULT_MODEL: &str = "claude-sonnet-4-6";
#[derive(Debug, Serialize, Deserialize)]
struct PersistedSession {
homeserver: String,
db_path: std::path::PathBuf,
user_session: MatrixSession,
#[serde(skip_serializing_if = "Option::is_none")]
sync_token: Option<String>,
}
#[derive(Clone, Debug)]
struct ChatMessage {
sender: OwnedUserId,
body: String,
is_self: bool,
/// Unix seconds. 0 if unknown.
ts: i64,
}
struct DaemonState {
own_user_id: OwnedUserId,
/// Per-room: the latest event_id that's been "shown" to Claude. Events
/// after this are "new" on the next invocation. Cleared on daemon restart.
last_shown: std::collections::HashMap<OwnedRoomId, OwnedEventId>,
pending_rooms: Vec<OwnedRoomId>,
rate_budget: u32,
rate_limit_per_min: u32,
last_rate_reset: std::time::Instant,
model: String,
}
const MAX_HISTORY: usize = 20;
const DEFAULT_RATE_LIMIT_PER_MIN: u32 = 1;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env())
.init();
tracing::info!("damocles-daemon starting");
let state_dir = paths::state_dir();
fs::create_dir_all(&state_dir).await?;
fs::create_dir_all(paths::identity_dir()).await?;
fs::create_dir_all(state_dir.join("rooms")).await?;
fs::create_dir_all(state_dir.join("people")).await?;
let session_file = paths::session_path();
let db_path = paths::db_path();
let config = load_config().await?;
let rate_limit_per_min = config
.rate_limit_per_min
.unwrap_or(DEFAULT_RATE_LIMIT_PER_MIN);
let model = config
.model
.clone()
.unwrap_or_else(|| DEFAULT_MODEL.to_owned());
let (client, sync_token) = if session_file.exists() {
restore_session(&session_file).await?
} else {
(login(&config, &db_path, &session_file).await?, None)
};
let own_user_id = client.user_id().context("not logged in")?.to_owned();
tracing::info!(
user = %own_user_id,
rate_limit = rate_limit_per_min,
model = %model,
"ready"
);
// Enable persistent event cache (matrix-sdk's sqlite store keeps the timeline)
client
.event_cache()
.subscribe()
.context("subscribe event cache")?;
let state = Arc::new(Mutex::new(DaemonState {
own_user_id,
last_shown: std::collections::HashMap::new(),
pending_rooms: Vec::new(),
rate_budget: rate_limit_per_min,
rate_limit_per_min,
last_rate_reset: std::time::Instant::now(),
model,
}));
let processor_state = state.clone();
let processor_client = client.clone();
tokio::spawn(async move {
process_loop(processor_state, processor_client).await;
});
sync(client, sync_token, &session_file, state).await
}
async fn load_config() -> anyhow::Result<Config> {
let data = fs::read_to_string(paths::config_path())
.await
.context("failed to read config.json")?;
serde_json::from_str(&data).context("failed to parse config.json")
}
async fn restore_session(session_file: &Path) -> anyhow::Result<(Client, Option<String>)> {
let data = fs::read_to_string(session_file).await?;
let session: PersistedSession = serde_json::from_str(&data)?;
tracing::info!(user = %session.user_session.meta.user_id, "restoring session");
let client = Client::builder()
.homeserver_url(&session.homeserver)
.sqlite_store(paths::db_path(), None)
.build()
.await?;
client.restore_session(session.user_session).await?;
Ok((client, session.sync_token))
}
async fn login(config: &Config, db_path: &Path, session_file: &Path) -> anyhow::Result<Client> {
tracing::info!(homeserver = %config.homeserver, user = %config.username, "logging in");
let client = Client::builder()
.homeserver_url(&config.homeserver)
.sqlite_store(db_path, None)
.build()
.await?;
client
.matrix_auth()
.login_username(&config.username, &config.password)
.initial_device_display_name("damocles-daemon")
.await?;
let user_session = client
.matrix_auth()
.session()
.context("no session after login")?;
let persisted = PersistedSession {
homeserver: config.homeserver.clone(),
db_path: db_path.to_owned(),
user_session,
sync_token: None,
};
fs::write(session_file, serde_json::to_string_pretty(&persisted)?).await?;
tracing::info!("session persisted");
Ok(client)
}
async fn sync(
client: Client,
initial_sync_token: Option<String>,
session_file: &Path,
state: Arc<Mutex<DaemonState>>,
) -> anyhow::Result<()> {
tracing::info!("initial sync (ignoring past messages)");
let filter = FilterDefinition::with_lazy_loading();
let mut sync_settings = SyncSettings::default().filter(filter.into());
if let Some(token) = initial_sync_token {
sync_settings = sync_settings.token(token);
}
loop {
match client.sync_once(sync_settings.clone()).await {
Ok(response) => {
sync_settings = sync_settings.token(response.next_batch.clone());
persist_sync_token(session_file, response.next_batch).await?;
break;
}
Err(e) => {
tracing::warn!("initial sync failed, retrying: {e}");
}
}
}
tracing::info!("synced, listening for messages");
client.add_event_handler(move |event: OriginalSyncRoomMessageEvent, room: Room| {
let state = state.clone();
async move {
on_room_message(event, room, state).await;
}
});
client.add_event_handler(on_stripped_state_member);
client.sync(sync_settings).await?;
bail!("sync loop exited unexpectedly")
}
async fn on_stripped_state_member(event: StrippedRoomMemberEvent, client: Client, room: Room) {
let Some(my_id) = client.user_id() else {
return;
};
if event.state_key != my_id {
return;
}
let room_id = room.room_id().to_owned();
tokio::spawn(async move {
tracing::info!(room = %room_id, "auto-joining invite");
let mut delay = 2u64;
loop {
match room.join().await {
Ok(()) => {
tracing::info!(room = %room_id, "joined");
if let Err(e) = ensure_room_notes(&room).await {
tracing::warn!(room = %room_id, "failed to write room notes: {e}");
}
return;
}
Err(e) => {
tracing::warn!(room = %room_id, "join failed, retry in {delay}s: {e}");
tokio::time::sleep(std::time::Duration::from_secs(delay)).await;
delay = (delay * 2).min(300);
if delay >= 300 {
tracing::error!(room = %room_id, "giving up on auto-join");
return;
}
}
}
}
});
}
async fn persist_sync_token(session_file: &Path, sync_token: String) -> anyhow::Result<()> {
let data = fs::read_to_string(session_file).await?;
let mut session: PersistedSession = serde_json::from_str(&data)?;
session.sync_token = Some(sync_token);
fs::write(session_file, serde_json::to_string_pretty(&session)?).await?;
Ok(())
}
async fn on_room_message(
event: OriginalSyncRoomMessageEvent,
room: Room,
state: Arc<Mutex<DaemonState>>,
) {
if room.state() != RoomState::Joined {
return;
}
let MessageType::Text(text_content) = &event.content.msgtype else {
return;
};
let room_id = room.room_id().to_owned();
let is_self = {
let state = state.lock().await;
event.sender == state.own_user_id
};
tracing::info!(
room = %room_id,
sender = %event.sender,
self_msg = is_self,
"{}",
text_content.body
);
if let Err(e) = ensure_room_notes(&room).await {
tracing::warn!(room = %room_id, "failed to ensure room notes: {e}");
}
if !is_self {
if let Err(e) = ensure_person_notes(&room, &event.sender).await {
tracing::warn!(sender = %event.sender, "failed to ensure person notes: {e}");
}
}
if !is_self {
let mut state = state.lock().await;
if !state.pending_rooms.contains(&room_id) {
state.pending_rooms.push(room_id);
}
}
}
/// Create state/rooms/<room_id>/notes.md if it doesn't exist.
async fn ensure_room_notes(room: &Room) -> anyhow::Result<()> {
let room_id = room.room_id();
let dir = paths::state_dir().join("rooms").join(room_id.as_str());
let file = dir.join("notes.md");
if file.exists() {
return Ok(());
}
fs::create_dir_all(&dir).await?;
let display_name = room
.display_name()
.await
.map(|n| n.to_string())
.unwrap_or_else(|_| room_id.to_string());
let now = chrono_now();
let body = format!("# {room_id}\n\nDisplay name: {display_name}\n\nFirst joined: {now}\n",);
fs::write(&file, body).await?;
tracing::info!(room = %room_id, "created room notes");
Ok(())
}
/// Create state/people/<user_id>/notes.md if it doesn't exist.
/// Pre-fill with display name (from this room's member info) and the room as
/// "first met in".
async fn ensure_person_notes(room: &Room, user_id: &UserId) -> anyhow::Result<()> {
let dir = paths::state_dir().join("people").join(user_id.as_str());
let file = dir.join("notes.md");
if file.exists() {
return Ok(());
}
fs::create_dir_all(&dir).await?;
let display_name = room
.get_member_no_sync(user_id)
.await
.ok()
.flatten()
.and_then(|m| m.display_name().map(ToOwned::to_owned))
.unwrap_or_default();
let display_line = if display_name.is_empty() {
String::new()
} else {
format!("Display name: {display_name}\n")
};
let room_id = room.room_id();
let now = chrono_now();
let body = format!("# {user_id}\n\n{display_line}First met in: {room_id} on {now}\n",);
fs::write(&file, body).await?;
tracing::info!(user = %user_id, "created person notes");
Ok(())
}
fn chrono_now() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0);
// Simple ISO-ish without pulling chrono. YYYY-MM-DD only is fine for notes.
let days = secs / 86400;
let (y, m, d) = days_to_ymd(days);
format!("{y:04}-{m:02}-{d:02}")
}
/// Format a unix-seconds timestamp as `YYYY-MM-DD HH:MM` UTC. Returns "?" for 0.
fn format_ts(secs: i64) -> String {
if secs == 0 {
return "?".into();
}
let days = secs.div_euclid(86400);
let day_secs = secs.rem_euclid(86400);
let (y, m, d) = days_to_ymd(days);
let h = day_secs / 3600;
let min = (day_secs % 3600) / 60;
format!("{y:04}-{m:02}-{d:02} {h:02}:{min:02}")
}
/// Convert days-since-1970-01-01 to (year, month, day). Civil-date algorithm.
fn days_to_ymd(z: i64) -> (i64, u32, u32) {
let z = z + 719_468;
let era = z.div_euclid(146_097);
let doe = z.rem_euclid(146_097) as u32;
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146_096) / 365;
let y = yoe as i64 + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153;
let d = doy - (153 * mp + 2) / 5 + 1;
let m = if mp < 10 { mp + 3 } else { mp - 9 };
(if m <= 2 { y + 1 } else { y }, m, d)
}
async fn process_loop(state: Arc<Mutex<DaemonState>>, client: Client) {
loop {
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
let room_id = {
let mut state = state.lock().await;
if state.last_rate_reset.elapsed() >= std::time::Duration::from_secs(60) {
state.rate_budget = state.rate_limit_per_min;
state.last_rate_reset = std::time::Instant::now();
}
if state.rate_budget == 0 {
continue;
}
state.pending_rooms.pop()
};
let Some(room_id) = room_id else {
continue;
};
let Some(room) = client.get_room(&room_id) else {
tracing::warn!(room = %room_id, "room not found in client");
continue;
};
// Snapshot last_shown for this room so we can mark seen vs new.
// If we don't have one in memory (e.g. fresh daemon start), seed from
// the stored read receipt so we don't reprocess old messages.
let in_memory = {
let s = state.lock().await;
s.last_shown.get(&room_id).cloned()
};
let prev_last_shown = if let Some(eid) = in_memory {
Some(eid)
} else {
let from_receipt = match room
.load_user_receipt(
matrix_sdk::ruma::events::receipt::ReceiptType::Read,
ReceiptThread::Unthreaded,
client.user_id().expect("logged in"),
)
.await
{
Ok(Some((eid, _))) => Some(eid),
Ok(None) => None,
Err(e) => {
tracing::warn!(room = %room_id, "failed to load receipt: {e}");
None
}
};
if let Some(ref eid) = from_receipt {
let mut s = state.lock().await;
s.last_shown.insert(room_id.clone(), eid.clone());
}
from_receipt
};
let room_name = room
.display_name()
.await
.map_or_else(|_| room_id.to_string(), |n| n.to_string());
// Load recent history from matrix-sdk's persistent event cache
let history = match load_recent_messages(&room, MAX_HISTORY).await {
Ok(h) => h,
Err(e) => {
tracing::error!(room = %room_id, "failed to load history: {e}");
continue;
}
};
let (own_user, model) = {
let state = state.lock().await;
(state.own_user_id.clone(), state.model.clone())
};
let chat_msgs: Vec<ChatMessage> = history
.iter()
.map(|(_, sender, body, ts)| ChatMessage {
sender: sender.clone(),
body: body.clone(),
is_self: sender == &own_user,
ts: *ts,
})
.collect();
// Determine seen split: everything before (and including) prev_last_shown is "seen"
let seen_idx = prev_last_shown
.as_ref()
.and_then(|id| history.iter().position(|(eid, _, _, _)| eid == id))
.map_or(0, |pos| pos + 1);
let new_last_event_id = history.last().map(|(eid, _, _, _)| eid.clone());
match invoke_claude(&room_id, &room_name, &chat_msgs, seen_idx, &model).await {
Ok(Some(response)) => {
if let Some(target_room) = client.get_room(&response.room) {
let content = RoomMessageEventContent::text_plain(&response.body);
match target_room.send(content).await {
Ok(_) => {
let mut state = state.lock().await;
state.rate_budget = state.rate_budget.saturating_sub(1);
if let Some(eid) = new_last_event_id.clone() {
state.last_shown.insert(room_id.clone(), eid);
}
tracing::info!(
room = %response.room,
"sent response ({} budget remaining)",
state.rate_budget
);
drop(state);
send_read_receipt(&room, new_last_event_id.clone()).await;
}
Err(e) => tracing::error!("failed to send: {e}"),
}
} else {
tracing::warn!(room = %response.room, "target room not found");
}
}
Ok(None) => {
tracing::debug!(room = %room_id, "claude chose to skip");
{
let mut state = state.lock().await;
if let Some(eid) = new_last_event_id.clone() {
state.last_shown.insert(room_id.clone(), eid);
}
}
send_read_receipt(&room, new_last_event_id.clone()).await;
}
Err(e) => {
tracing::error!(room = %room_id, "claude invocation failed: {e}");
}
}
}
}
async fn send_read_receipt(room: &Room, event_id: Option<OwnedEventId>) {
let Some(eid) = event_id else {
return;
};
if let Err(e) = room
.send_single_receipt(CreateReceiptType::Read, ReceiptThread::Unthreaded, eid)
.await
{
tracing::warn!(room = %room.room_id(), "failed to send read receipt: {e}");
}
}
/// Load the last N text messages from the room's persistent event cache.
/// Returns oldest-first list of (event_id, sender, body, ts_secs).
async fn load_recent_messages(
room: &Room,
limit: usize,
) -> anyhow::Result<Vec<(OwnedEventId, OwnedUserId, String, i64)>> {
use matrix_sdk::ruma::events::AnySyncTimelineEvent;
let (cache, _handles) = room.event_cache().await?;
let events = cache.events().await;
let mut out: Vec<(OwnedEventId, OwnedUserId, String, i64)> = Vec::new();
for ev in events.iter().rev() {
if out.len() >= limit {
break;
}
let raw = ev.raw();
let Ok(deserialized) = raw.deserialize() else {
continue;
};
if let AnySyncTimelineEvent::MessageLike(msg) = deserialized {
if let matrix_sdk::ruma::events::AnySyncMessageLikeEvent::RoomMessage(
matrix_sdk::ruma::events::SyncMessageLikeEvent::Original(orig),
) = msg
{
if let MessageType::Text(text) = &orig.content.msgtype {
let ts_ms: u64 = orig.origin_server_ts.0.into();
let ts_secs: i64 = i64::try_from(ts_ms).unwrap_or(0) / 1000;
out.push((
orig.event_id.clone(),
orig.sender.clone(),
text.body.clone(),
ts_secs,
));
}
}
}
}
out.reverse();
Ok(out)
}
struct ClaudeResponse {
room: OwnedRoomId,
body: String,
}
async fn invoke_claude(
source_room: &OwnedRoomId,
room_name: &str,
history: &[ChatMessage],
seen_idx: usize,
model: &str,
) -> anyhow::Result<Option<ClaudeResponse>> {
let identity_dir = paths::identity_dir();
let identity_str = identity_dir.to_string_lossy();
let mut prompt = String::new();
writeln!(prompt, "[room_id: {source_room}]").unwrap();
writeln!(prompt, "[room_name: {room_name}]").unwrap();
writeln!(
prompt,
"[room notes path: ../rooms/{source_room}/notes.md (create dir if needed)]"
)
.unwrap();
// Collect unique non-self senders for per-person note discovery
let mut senders: Vec<&OwnedUserId> = history
.iter()
.filter(|m| !m.is_self)
.map(|m| &m.sender)
.collect();
senders.sort();
senders.dedup();
if !senders.is_empty() {
writeln!(
prompt,
"\n[people in this room — check ../people/<user_id>/notes.md for each]"
)
.unwrap();
for s in &senders {
writeln!(prompt, " {s}").unwrap();
}
}
let seen = seen_idx.min(history.len());
let (old, new) = history.split_at(seen);
if !old.is_empty() {
writeln!(prompt, "\n[previously seen messages — for context]").unwrap();
for msg in old {
let prefix = if msg.is_self { "(you) " } else { "" };
let ts = format_ts(msg.ts);
writeln!(prompt, "[{ts}] {prefix}{}: {}", msg.sender, msg.body).unwrap();
}
}
writeln!(prompt, "\n[new messages — respond to these]").unwrap();
if new.is_empty() {
writeln!(prompt, "(none)").unwrap();
} else {
for msg in new {
let prefix = if msg.is_self { "(you) " } else { "" };
let ts = format_ts(msg.ts);
writeln!(prompt, "[{ts}] {prefix}{}: {}", msg.sender, msg.body).unwrap();
}
}
tracing::info!("invoking claude: {} new, {} seen", new.len(), old.len());
tracing::trace!("full prompt:\n{prompt}");
use tokio::process::Command;
let mut cmd = Command::new("claude");
cmd.args([
"--print",
"--model",
model,
"--add-dir",
&identity_str,
"--allowedTools",
"Read Edit Write Glob Grep",
"-p",
&prompt,
]);
cmd.current_dir(&identity_dir);
cmd.stdin(std::process::Stdio::null());
let output = cmd.output().await.context("failed to run claude")?;
let stderr = String::from_utf8_lossy(&output.stderr);
let stdout = String::from_utf8_lossy(&output.stdout);
if !output.status.success() {
bail!(
"claude exited with {}:\nstdout: {}\nstderr: {}",
output.status,
stdout,
stderr
);
}
if !stderr.is_empty() {
tracing::warn!("claude stderr: {stderr}");
}
let raw = String::from_utf8_lossy(&output.stdout).to_string();
Ok(parse_response(&raw, source_room))
}
fn parse_response(raw: &str, default_room: &OwnedRoomId) -> Option<ClaudeResponse> {
let trimmed = raw.trim();
if trimmed.starts_with("---") {
let parts: Vec<&str> = trimmed.splitn(3, "---").collect();
if parts.len() >= 3 {
let frontmatter = parts[1].trim();
let body = parts[2].trim();
if frontmatter.contains("skip: true") || frontmatter.contains("skip:true") {
return None;
}
let room = frontmatter
.lines()
.find(|l| l.starts_with("room:"))
.and_then(|l| l.strip_prefix("room:"))
.and_then(|r| r.trim().parse().ok())
.unwrap_or_else(|| default_room.clone());
if body.is_empty() {
return None;
}
return Some(ClaudeResponse {
room,
body: body.to_owned(),
});
}
}
if trimmed.is_empty() {
return None;
}
Some(ClaudeResponse {
room: default_room.clone(),
body: trimmed.to_owned(),
})
}
#[cfg(test)]
mod tests {
use super::*;
fn test_room() -> OwnedRoomId {
"!test:example.com".parse().unwrap()
}
#[test]
fn parse_frontmatter_response() {
let raw = "---\nroom: !other:server\n---\nhello world";
let resp = parse_response(raw, &test_room()).unwrap();
assert_eq!(resp.room.as_str(), "!other:server");
assert_eq!(resp.body, "hello world");
}
#[test]
fn parse_skip_response() {
let raw = "---\nskip: true\n---\n";
assert!(parse_response(raw, &test_room()).is_none());
}
#[test]
fn parse_plain_response() {
let raw = "just a message";
let resp = parse_response(raw, &test_room()).unwrap();
assert_eq!(resp.room, test_room());
assert_eq!(resp.body, "just a message");
}
#[test]
fn parse_empty_response() {
assert!(parse_response("", &test_room()).is_none());
assert!(parse_response(" \n ", &test_room()).is_none());
}
#[test]
fn parse_default_room() {
let raw = "---\n---\nhello";
let resp = parse_response(raw, &test_room()).unwrap();
assert_eq!(resp.room, test_room());
}
}