@@ -7,14 +7,13 @@ use anyhow::{anyhow, Context, Result}; use klbr_core::{ harness_block::{FormatStyle, DEFAULT_FORMAT_STYLE}, interrupt::{ExternalEvent, Interrupt}, - models::ToolDef, tools::{Subroutines, Tool, ToolFuture}, AgentEvent, }; use rusqlite::{params, Connection, OptionalExtension, Row}; -use serde_json::{json, Value}; +use serde_json::Value; use tokio::{ sync::{broadcast, mpsc}, task::JoinHandle, }; @@ -28,8 +27,23 @@ }, util::Timestamp, }; +mod formatting; +mod inbox; +mod tool_defs; + +use formatting::{ + format_discord_batch, format_discord_batch_time, format_discord_timestamp, format_history, +}; +#[cfg(test)] +use formatting::{format_discord_timestamp_secs, one_line}; +use inbox::{DiscordInboxItem, DiscordInboxStore}; +use tool_defs::{ + discord_fetch_history_def, discord_list_conversations_def, discord_mark_def, discord_react_def, + discord_send_def, +}; + #[derive(Clone, Debug)] pub struct DiscordConfig { pub token: String, pub source: String, @@ -81,26 +95,8 @@ is_dm: bool, author_is_bot: bool, } -#[derive(Clone, Debug)] -struct DiscordInboxStore { - conn: Arc>, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -struct DiscordInboxItem { - item_id: String, - conversation_id: String, - channel_id: String, - message_id: String, - author_name: String, - content: String, - bucket: String, - status: String, - timestamp: i64, -} - impl DiscordConfig { pub fn from_env() -> Result> { let enabled = env_bool("KLBR_DISCORD_ENABLED"); if enabled == Some(false) { @@ -142,185 +138,8 @@ })) } } -impl DiscordInboxStore { - fn open(path: impl AsRef) -> Result { - let conn = Connection::open(path)?; - let store = Self { - conn: Arc::new(Mutex::new(conn)), - }; - store.init_schema()?; - Ok(store) - } - - fn init_schema(&self) -> Result<()> { - self.conn.lock().unwrap().execute_batch( - "CREATE TABLE IF NOT EXISTS discord_inbox_items ( - item_id TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - channel_id TEXT NOT NULL, - message_id TEXT NOT NULL UNIQUE, - author_name TEXT NOT NULL, - content TEXT NOT NULL, - bucket TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - action TEXT NOT NULL DEFAULT 'none', - note TEXT, - message_ts INTEGER NOT NULL, - first_seen_ts INTEGER NOT NULL DEFAULT (unixepoch()), - last_seen_ts INTEGER NOT NULL DEFAULT (unixepoch()), - decision_ts INTEGER - ); - CREATE INDEX IF NOT EXISTS idx_discord_inbox_status - ON discord_inbox_items(status, last_seen_ts DESC); - CREATE INDEX IF NOT EXISTS idx_discord_inbox_channel - ON discord_inbox_items(channel_id, status, last_seen_ts DESC);", - )?; - Ok(()) - } - - fn upsert_pending(&self, message: &DiscordIncomingMessage) -> Result> { - let Some(bucket) = pending_bucket(message) else { - return Ok(None); - }; - - let conn = self.conn.lock().unwrap(); - conn.execute( - "INSERT INTO discord_inbox_items ( - item_id, conversation_id, channel_id, message_id, author_name, - content, bucket, status, action, message_ts - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, 'pending', 'none', ?8) - ON CONFLICT(item_id) DO UPDATE SET - conversation_id = excluded.conversation_id, - channel_id = excluded.channel_id, - author_name = excluded.author_name, - content = excluded.content, - bucket = excluded.bucket, - last_seen_ts = unixepoch()", - params![ - message.message_id, - message.conversation_id, - message.channel_id, - message.message_id, - message.author_name, - message.content, - bucket, - message.timestamp, - ], - )?; - Ok(Some(message.message_id.clone())) - } - - fn pending_items(&self, limit: usize) -> Result> { - let conn = self.conn.lock().unwrap(); - let mut stmt = conn.prepare( - "SELECT item_id, conversation_id, channel_id, message_id, author_name, - content, bucket, status, message_ts - FROM discord_inbox_items - WHERE status = 'pending' - ORDER BY message_ts DESC, last_seen_ts DESC - LIMIT ?1", - )?; - let rows = stmt.query_map(params![limit as i64], inbox_item_from_row)?; - rows.collect::, _>>() - .map_err(Into::into) - } - - fn demote_stale_pending(&self, older_than_secs: i64) -> Result { - let changed = self.conn.lock().unwrap().execute( - "UPDATE discord_inbox_items - SET status = 'seen', - action = CASE WHEN action = 'none' THEN 'noted' ELSE action END, - decision_ts = unixepoch(), - last_seen_ts = unixepoch() - WHERE status = 'pending' - AND message_ts <= unixepoch() - ?1", - params![older_than_secs], - )?; - Ok(changed) - } - - fn pending_item_for_message(&self, message_id: &str) -> Result> { - self.conn - .lock() - .unwrap() - .query_row( - "SELECT item_id - FROM discord_inbox_items - WHERE message_id = ?1 AND status = 'pending' - LIMIT 1", - params![message_id], - |row| row.get(0), - ) - .optional() - .map_err(Into::into) - } - - fn conversation_summaries(&self) -> Result> { - // returns (conversation_id, channel_id, bucket, latest_ts) per conversation - let conn = self.conn.lock().unwrap(); - let mut stmt = conn.prepare( - "SELECT conversation_id, channel_id, - CASE WHEN SUM(bucket = 'dm') > 0 THEN 'dm' ELSE 'guild' END AS bucket, - MAX(message_ts) AS latest_ts - FROM discord_inbox_items - GROUP BY conversation_id - ORDER BY latest_ts DESC - LIMIT 50", - )?; - let rows = stmt.query_map([], |row| { - Ok(( - row.get::<_, String>(0)?, - row.get::<_, String>(1)?, - row.get::<_, String>(2)?, - row.get::<_, i64>(3)?, - )) - })?; - rows.collect::, _>>() - .map_err(Into::into) - } - - fn pending_dm_item_for_channel(&self, channel_id: &str) -> Result> { - self.conn - .lock() - .unwrap() - .query_row( - "SELECT item_id - FROM discord_inbox_items - WHERE channel_id = ?1 - AND bucket = 'dm' - AND status = 'pending' - ORDER BY message_ts DESC - LIMIT 1", - params![channel_id], - |row| row.get(0), - ) - .optional() - .map_err(Into::into) - } - - fn mark(&self, item_id: &str, status: &str, action: &str, note: Option<&str>) -> Result<()> { - validate_inbox_status(status)?; - validate_inbox_action(action)?; - let item_id = normalize_item_id(item_id); - let changed = self.conn.lock().unwrap().execute( - "UPDATE discord_inbox_items - SET status = ?1, - action = ?2, - note = ?3, - decision_ts = unixepoch(), - last_seen_ts = unixepoch() - WHERE item_id = ?4", - params![status, action, note, item_id], - )?; - if changed == 0 { - anyhow::bail!("unknown Discord inbox item {item_id:?}"); - } - Ok(()) - } -} - impl DiscordRuntime { pub async fn new(config: DiscordConfig, db_path: impl AsRef) -> Result { let _ = rustls::crypto::ring::default_provider().install_default(); @@ -993,130 +812,8 @@ self.inbox.pending_dm_item_for_channel(channel_id) } } -fn discord_send_def() -> ToolDef { - ToolDef::function( - "discord_send", - "send a Discord message. use this when a Discord response is useful; many ambient batches need no reply. plain assistant text only goes to the local chat. conversation_id is optional when the current Discord batch has a single conversation. pass source_item_id when responding to a pending item; successful sends mark that item acted, so do not call discord_mark afterward.", - json!({ - "type": "object", - "properties": { - "conversation_id": { - "type": "string", - "description": "Discord conversation id from the batch, like d1 or conv:d1; omit only for a single-conversation batch" - }, - "content": { - "type": "string", - "description": "message content to send, max 2000 characters" - }, - "source_item_id": { - "type": "string", - "description": "optional pending inbox item id shown as msg:; marks it acted after sending" - } - }, - "required": ["content"] - }), - ) -} - -fn discord_fetch_history_def() -> ToolDef { - ToolDef::function( - "discord_fetch_history", - "fetch recent Discord message history for a conversation. use this when the current batch lacks enough context. conversation_id is optional when the current batch has a single conversation.", - json!({ - "type": "object", - "properties": { - "conversation_id": { - "type": "string", - "description": "Discord conversation id from the batch, like d1 or conv:d1; omit only for a single-conversation batch" - }, - "before_message_id": { - "type": "string", - "description": "optional msg: from the batch; fetch messages before it" - }, - "after_message_id": { - "type": "string", - "description": "optional msg: from the batch; fetch messages after it" - }, - "limit": { - "type": "integer", - "description": "number of messages to fetch, 1-100; default 20" - } - }, - "required": [] - }), - ) -} - -fn discord_react_def() -> ToolDef { - ToolDef::function( - "discord_react", - "add a reaction to a Discord message. prefer real unicode emoji like 😍, 👍, or ❤️. only use a custom emoji string like <:name:id> when the exact string came from Discord; never invent custom emoji ids. conversation_id is optional when the current Discord batch has a single conversation.", - json!({ - "type": "object", - "properties": { - "conversation_id": { - "type": "string", - "description": "Discord conversation id from the batch, like d1 or conv:d1; omit only for a single-conversation batch" - }, - "message_id": { - "type": "string", - "description": "msg: of the message to react to, as shown in the batch" - }, - "emoji": { - "type": "string", - "description": "unicode emoji is preferred; custom emoji must be an exact Discord form <:name:id> / , copied from Discord rather than guessed" - } - }, - "required": ["message_id", "emoji"] - }), - ) -} - -fn discord_list_conversations_def() -> ToolDef { - ToolDef::function( - "discord_list_conversations", - "list known Discord conversations with their ids, type (dm/guild), and last activity. use this to find a conversation_id when it is not in the current batch.", - json!({ - "type": "object", - "properties": {}, - "required": [] - }), - ) -} - -fn discord_mark_def() -> ToolDef { - ToolDef::function( - "discord_mark", - "mark a pending Discord inbox item as seen, acted, or ignored so future batches remember the decision. use this when you are not sending a Discord response, especially ignored when you deliberately choose not to respond. successful discord_send already marks the item acted.", - json!({ - "type": "object", - "properties": { - "item_id": { - "type": "string", - "description": "pending Discord inbox message id, either raw snowflake or the displayed msg:" - }, - "status": { - "type": "string", - "enum": ["pending", "seen", "acted", "ignored"], - "description": "new item status" - }, - "action": { - "type": "string", - "enum": ["none", "replied", "messaged", "dismissed", "noted"], - "description": "optional decision/action label" - }, - "note": { - "type": "string", - "description": "optional short decision note" - } - }, - "required": ["item_id", "status"] - }), - ) -} - fn env_bool(name: &str) -> Option { std::env::var(name).ok().map(|value| { matches!( value.to_ascii_lowercase().as_str(), @@ -1247,40 +944,8 @@ .or_else(|| s.strip_prefix("item:")) .unwrap_or(s) } -fn inbox_item_from_row(row: &Row<'_>) -> rusqlite::Result { - Ok(DiscordInboxItem { - item_id: row.get(0)?, - conversation_id: row.get(1)?, - channel_id: row.get(2)?, - message_id: row.get(3)?, - author_name: row.get(4)?, - content: row.get(5)?, - bucket: row.get(6)?, - status: row.get(7)?, - timestamp: row.get(8)?, - }) -} - -fn validate_inbox_status(status: &str) -> Result<()> { - match status { - "pending" | "seen" | "acted" | "ignored" => Ok(()), - other => anyhow::bail!( - "invalid Discord inbox status {other:?}; expected pending, seen, acted, or ignored" - ), - } -} - -fn validate_inbox_action(action: &str) -> Result<()> { - match action { - "none" | "replied" | "messaged" | "dismissed" | "noted" => Ok(()), - other => anyhow::bail!( - "invalid Discord inbox action {other:?}; expected none, replied, messaged, dismissed, or noted" - ), - } -} - fn resolve_channel_arg_from_state( state: &DiscordState, args: &Value, allow_active: bool, @@ -1424,260 +1089,8 @@ } Some(name.replace('-', "_").to_ascii_lowercase()) } -fn format_history( - channel_id: Id, - conversation_id: Option<&str>, - messages: &[Message], - style: FormatStyle, -) -> String { - match style { - FormatStyle::Brackets => { - let conversation = conversation_id - .map(|id| format!("conv:{id}")) - .unwrap_or_else(|| format!("channel:{channel_id}")); - let mut lines = vec![ - "[discord history]".to_string(), - conversation, - "messages:".to_string(), - ]; - - for message in messages.iter().rev() { - let author = message - .author - .global_name - .as_deref() - .unwrap_or(&message.author.name); - lines.push(format!( - "[msg:{} {}] {}: {}", - message.id, - format_discord_timestamp(message.timestamp), - author, - one_line(&message.content) - )); - } - - lines.join("\n") - } - FormatStyle::Xml => { - let mut lines = vec![]; - let conversation_attr = conversation_id - .map(|id| format!(" conversation_id=\"{id}\"")) - .unwrap_or_else(|| format!(" channel_id=\"{channel_id}\"")); - lines.push(format!("", conversation_attr)); - - for message in messages.iter().rev() { - let author = message - .author - .global_name - .as_deref() - .unwrap_or(&message.author.name); - lines.push(format!( - " {}", - message.id, - format_discord_timestamp(message.timestamp), - author, - one_line(&message.content) - )); - } - - lines.push("".to_string()); - lines.join("\n") - } - } -} - -fn format_discord_batch( - messages: &[DiscordIncomingMessage], - pending: &[DiscordInboxItem], - include_instructions: bool, - style: FormatStyle, -) -> String { - // pending items mapped by message_id - let batch_pending: HashMap<&str, &DiscordInboxItem> = pending - .iter() - .map(|item| (item.message_id.as_str(), item)) - .collect(); - - match style { - FormatStyle::Brackets => { - let mut lines = vec!["[discord message batch]".to_string()]; - if include_instructions { - lines.push("messages are context, not direct requests".to_string()); - lines.push("for dm: reply if a response is needed".to_string()); - lines.push( - "for guild: channel context only — replying is optional, use judgment" - .to_string(), - ); - lines.push("use discord_send to respond (marks pending items acted); use discord_mark only when not sending".to_string()); - } - lines.push("messages:".to_string()); - - for message in messages { - let tag = batch_pending - .get(message.message_id.as_str()) - .map(|item| item.bucket.as_str()) - .unwrap_or(if message.is_dm { "dm" } else { "guild" }); - lines.push(format!( - "[conv:{} msg:{} {} {}] {}: {}", - message.conversation_id, - message.message_id, - tag, - format_discord_batch_time(message.timestamp), - message.author_name, - one_line(&message.content) - )); - } - - let batch_msg_ids: HashSet<&str> = - messages.iter().map(|m| m.message_id.as_str()).collect(); - let older_pending: Vec<&DiscordInboxItem> = pending - .iter() - .filter(|item| !batch_msg_ids.contains(item.message_id.as_str())) - .collect(); - - if !older_pending.is_empty() { - lines.push("pending:".to_string()); - for item in older_pending { - lines.push(format!( - "[msg:{} {} conv:{} {}] {}: {}", - item.item_id, - item.bucket, - item.conversation_id, - format_discord_batch_time(item.timestamp), - item.author_name, - one_line(&item.content) - )); - } - } - - lines.join("\n") - } - FormatStyle::Xml => { - let mut lines = vec!["".to_string()]; - if include_instructions { - lines.push(" ".to_string()); - lines.push(" messages are context, not direct requests".to_string()); - lines.push(" for dm: reply if a response is needed".to_string()); - lines.push( - " for guild: channel context only — replying is optional, use judgment" - .to_string(), - ); - lines.push(" use discord_send to respond (marks pending items acted); use discord_mark only when not sending".to_string()); - lines.push(" ".to_string()); - } - - lines.push(" ".to_string()); - for message in messages { - let pending_item = batch_pending.get(message.message_id.as_str()); - let tag = pending_item - .map(|item| item.bucket.as_str()) - .unwrap_or(if message.is_dm { "dm" } else { "guild" }); - - let item_attr = pending_item - .map(|item| format!(" item_id=\"{}\"", item.item_id)) - .unwrap_or_default(); - - lines.push(format!( - " {}", - message.message_id, - item_attr, - message.conversation_id, - tag, - format_discord_batch_time(message.timestamp), - message.author_name, - one_line(&message.content) - )); - } - lines.push(" ".to_string()); - - let batch_msg_ids: HashSet<&str> = - messages.iter().map(|m| m.message_id.as_str()).collect(); - let older_pending: Vec<&DiscordInboxItem> = pending - .iter() - .filter(|item| !batch_msg_ids.contains(item.message_id.as_str())) - .collect(); - - if !older_pending.is_empty() { - lines.push(" ".to_string()); - for item in older_pending { - lines.push(format!( - " {}", - item.message_id, - item.item_id, - item.conversation_id, - item.bucket, - format_discord_batch_time(item.timestamp), - item.author_name, - one_line(&item.content) - )); - } - lines.push(" ".to_string()); - } - - lines.push("".to_string()); - lines.join("\n") - } - } -} - -fn format_discord_timestamp(timestamp: Timestamp) -> String { - format_discord_timestamp_secs(timestamp.as_secs()) -} - -fn format_discord_timestamp_secs(secs: i64) -> String { - let Some(datetime) = chrono::DateTime::from_timestamp(secs, 0) else { - return secs.to_string(); - }; - datetime.format("%Y-%m-%d %H:%M:%S").to_string() -} - -fn format_discord_batch_time(secs: i64) -> String { - let Some(datetime) = chrono::DateTime::from_timestamp(secs, 0) else { - return secs.to_string(); - }; - datetime.format("%Y-%m-%d %H:%M:%S").to_string() -} - -fn one_line(text: &str) -> String { - use std::sync::OnceLock; - static DATA_URL_RE: OnceLock = OnceLock::new(); - let re = DATA_URL_RE.get_or_init(|| { - regex::Regex::new(r"data:([a-zA-Z0-9\-+\.]+/[a-zA-Z0-9\-+\.]+);base64,([a-zA-Z0-9/+=]+)") - .unwrap() - }); - - let mut media_urls = Vec::new(); - let mut cleaned_text = text.to_string(); - - while let Some(mat) = re.find(&cleaned_text) { - media_urls.push(mat.as_str().to_string()); - cleaned_text.replace_range(mat.range(), ""); - } - - let normalized = cleaned_text - .split_whitespace() - .collect::>() - .join(" "); - let mut out = if normalized.chars().count() <= 500 { - normalized - } else { - let mut truncated = normalized.chars().take(500).collect::(); - truncated.push_str("..."); - truncated - }; - - for url in media_urls { - if !out.is_empty() { - out.push(' '); - } - out.push_str(&url); - } - - out -} - #[cfg(test)] mod tests { use super::{ choose_reply_reference_from_state, conversation_id_for_channel_from_state, env_bool,