use std::{
collections::{HashMap, HashSet},
sync::{Arc, Mutex},
time::Duration,
};
use anyhow::{anyhow, Context, Result};
use serde_json::Value;
use tokio::{sync::mpsc, task::JoinHandle};
use twilight_gateway::{Event, EventTypeFlags, Intents, Shard, ShardId, StreamExt as _};
use twilight_http::{request::channel::reaction::RequestReactionType, Client as HttpClient};
use twilight_model::{
channel::{message::AllowedMentions, Message},
id::{
marker::{ChannelMarker, EmojiMarker, MessageMarker, UserMarker},
Id,
},
util::Timestamp,
};
mod formatting;
use formatting::{format_discord_batch, format_discord_timestamp, format_history};
#[derive(Clone, Debug)]
pub struct DiscordConfig {
pub token: String,
pub source: String,
pub event_mode: DiscordEventMode,
pub include_bot_messages: bool,
pub batch_interval: Duration,
pub channel_ids: HashSet<String>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DiscordEventMode {
Mentions,
All,
}
#[derive(Clone)]
pub struct DiscordRuntime {
config: DiscordConfig,
http: Arc<HttpClient>,
bot_user_id: Id<UserMarker>,
state: Arc<Mutex<DiscordState>>,
}
#[derive(Clone, Debug, Default)]
struct DiscordState {
latest_message_by_channel: HashMap<String, String>,
reply_target_by_channel: HashMap<String, String>,
conversation_by_channel: HashMap<String, String>,
channel_by_conversation: HashMap<String, String>,
active_channel_id: Option<String>,
next_conversation_id: usize,
}
#[derive(Clone, Debug)]
struct DiscordIncomingMessage {
conversation_id: String,
content: String,
author_id: String,
author_name: String,
message_id: String,
timestamp: i64,
guild_id: Option<String>,
channel_id: String,
mentions_bot: bool,
is_dm: bool,
author_is_bot: bool,
}
/// A normalized Discord batch for admission into the runtime-owned inbox.
/// The adapter owns transport parsing; the daemon owns durability and scheduling.
#[derive(Clone, Debug)]
pub struct DiscordIngress {
pub source: String,
pub conversation_id: String,
pub content: String,
pub message_ids: Vec<String>,
pub timestamp_ms: Option<i64>,
pub metadata: Value,
}
impl DiscordConfig {
pub fn from_env() -> Result<Option<Self>> {
let enabled = env_bool("KLBR_DISCORD_ENABLED");
if enabled == Some(false) {
return Ok(None);
}
let token = std::env::var("KLBR_DISCORD_TOKEN")
.or_else(|_| std::env::var("DISCORD_TOKEN"))
.ok();
let Some(token) = token else {
if enabled == Some(true) {
anyhow::bail!("KLBR_DISCORD_ENABLED is set but KLBR_DISCORD_TOKEN is missing");
}
return Ok(None);
};
let event_mode = match std::env::var("KLBR_DISCORD_EVENT_MODE")
.unwrap_or_else(|_| "all".to_string())
.to_ascii_lowercase()
.as_str()
{
"all" => DiscordEventMode::All,
"mentions" | "mention" | "" => DiscordEventMode::Mentions,
other => anyhow::bail!(
"invalid KLBR_DISCORD_EVENT_MODE={other:?}; expected 'mentions' or 'all'"
),
};
Ok(Some(Self {
token,
source: std::env::var("KLBR_DISCORD_SOURCE").unwrap_or_else(|_| "discord".to_string()),
event_mode,
include_bot_messages: env_bool("KLBR_DISCORD_INCLUDE_BOTS").unwrap_or(false),
batch_interval: env_duration_secs("KLBR_DISCORD_BATCH_SECS", 8)?,
channel_ids: env_channel_ids("KLBR_DISCORD_CHANNEL_IDS")?,
}))
}
}
impl DiscordRuntime {
pub async fn new(config: DiscordConfig) -> Result<Self> {
let _ = rustls::crypto::ring::default_provider().install_default();
let http = Arc::new(HttpClient::new(config.token.clone()));
let current_user = http
.current_user()
.await
.context("failed to fetch Discord current user")?
.model()
.await
.context("failed to decode Discord current user")?;
Ok(Self {
config,
http,
bot_user_id: current_user.id,
state: Arc::new(Mutex::new(DiscordState::default())),
})
}
pub async fn from_env() -> Result<Option<Arc<Self>>> {
let Some(config) = DiscordConfig::from_env()? else {
return Ok(None);
};
Ok(Some(Arc::new(Self::new(config).await?)))
}
pub fn spawn_gateway(
self: Arc<Self>,
ingress_tx: mpsc::Sender<DiscordIngress>,
) -> JoinHandle<()> {
tokio::spawn(async move {
if let Err(err) = self.run_gateway(ingress_tx).await {
tracing::error!(%err, "discord gateway stopped");
}
})
}
async fn run_gateway(self: Arc<Self>, ingress_tx: mpsc::Sender<DiscordIngress>) -> Result<()> {
let intents = Intents::GUILD_MESSAGES | Intents::DIRECT_MESSAGES | Intents::MESSAGE_CONTENT;
let mut shard = Shard::new(ShardId::ONE, self.config.token.clone(), intents);
let (batch_tx, batch_rx) = mpsc::channel(1024);
let batcher = self.clone().spawn_batcher(batch_rx, ingress_tx);
tracing::info!(
bot_user_id = %self.bot_user_id,
mode = ?self.config.event_mode,
batch_interval_ms = self.config.batch_interval.as_millis(),
filtered_channel_count = self.config.channel_ids.len(),
"discord gateway starting"
);
let result = async {
while let Some(item) = shard.next_event(EventTypeFlags::all()).await {
let event = match item {
Ok(event) => event,
Err(err) => {
tracing::warn!(?err, "discord gateway event error");
continue;
}
};
match event {
Event::Ready(ready) => {
tracing::info!(
bot_user_id = %ready.user.id,
username = %ready.user.name,
"discord gateway ready"
);
}
Event::MessageCreate(message) => {
if let Some(message) = self.incoming_message(&message).await {
batch_tx
.send(message)
.await
.context("discord batch channel closed")?;
}
}
_ => {}
}
}
Ok(())
}
.await;
batcher.abort();
result
}
fn spawn_batcher(
self: Arc<Self>,
mut batch_rx: mpsc::Receiver<DiscordIncomingMessage>,
ingress_tx: mpsc::Sender<DiscordIngress>,
) -> JoinHandle<()> {
tokio::spawn(async move {
let mut batch_count = 0usize;
while let Some(first) = batch_rx.recv().await {
let directed = first.is_dm || first.mentions_bot;
let mut batch = vec![first];
if !directed {
tokio::time::sleep(self.config.batch_interval).await;
}
while let Ok(message) = batch_rx.try_recv() {
batch.push(message);
}
let count = batch.len();
let include_instructions = batch_count.is_multiple_of(6) || count >= 5;
batch_count += 1;
self.note_batch_targets(&batch);
let ingress = self.batch_ingress(&batch, include_instructions);
if let Err(err) = ingress_tx.send(ingress).await {
tracing::warn!(?err, "runtime ingress channel closed");
break;
}
tracing::debug!(
message_count = count,
directed,
"admitted Discord message batch to runtime transport"
);
}
})
}
fn batch_ingress(
&self,
batch: &[DiscordIncomingMessage],
include_instructions: bool,
) -> DiscordIngress {
let first = batch.first().expect("discord batch must not be empty");
let all_same_conversation = batch
.iter()
.all(|message| message.conversation_id == first.conversation_id);
let conversation_id = if all_same_conversation {
first.conversation_id.clone()
} else {
format!("batch:{}:messages", self.config.source)
};
let source = if batch.iter().any(|message| message.is_dm) {
"discord.dm"
} else if batch.iter().any(|message| message.mentions_bot) {
"discord.mention"
} else {
"discord.ambient"
};
let metadata = serde_json::json!({
"transport": self.config.source,
"messages": batch.iter().map(|message| serde_json::json!({
"message_id": message.message_id,
"channel_id": message.channel_id,
"guild_id": message.guild_id,
"author_id": message.author_id,
"author_name": message.author_name,
"author_is_bot": message.author_is_bot,
"is_dm": message.is_dm,
"mentions_bot": message.mentions_bot,
})).collect::<Vec<_>>(),
});
DiscordIngress {
source: source.into(),
conversation_id,
content: format_discord_batch(batch, include_instructions),
message_ids: batch
.iter()
.map(|message| message.message_id.clone())
.collect(),
timestamp_ms: batch
.iter()
.map(|message| message.timestamp.saturating_mul(1_000))
.min(),
metadata,
}
}
fn note_seen_message(&self, channel_id: &str, message_id: &str) {
let mut state = self.state.lock().expect("discord runtime state poisoned");
state
.latest_message_by_channel
.insert(channel_id.to_string(), message_id.to_string());
}
fn note_batch_targets(&self, batch: &[DiscordIncomingMessage]) {
let mut latest_by_channel: HashMap<&str, &str> = HashMap::new();
for message in batch {
latest_by_channel.insert(&message.channel_id, &message.message_id);
}
let mut state = self.state.lock().expect("discord runtime state poisoned");
note_batch_targets_in_state(&mut state, latest_by_channel);
}
fn choose_reply_reference(&self, channel_id: &str) -> Option<String> {
let mut state = self.state.lock().expect("discord runtime state poisoned");
choose_reply_reference_from_state(&mut state, channel_id)
}
fn conversation_id_for_channel(&self, channel_id: &str) -> String {
let mut state = self.state.lock().expect("discord runtime state poisoned");
conversation_id_for_channel_from_state(&mut state, channel_id)
}
fn existing_conversation_id_for_channel(&self, channel_id: &str) -> Option<String> {
let state = self.state.lock().expect("discord runtime state poisoned");
state.conversation_by_channel.get(channel_id).cloned()
}
fn resolve_channel(&self, args: &Value, allow_active: bool) -> Result<Id<ChannelMarker>> {
let state = self.state.lock().expect("discord runtime state poisoned");
let channel_id = resolve_channel_arg_from_state(&state, args, allow_active)?;
parse_discord_id(&channel_id, "Discord channel")
}
fn resolve_send_channel(&self, args: &Value) -> Result<Id<ChannelMarker>> {
self.resolve_channel(args, true)
}
async fn incoming_message(
&self,
message: &twilight_model::gateway::payload::incoming::MessageCreate,
) -> Option<DiscordIncomingMessage> {
if message.author.id == self.bot_user_id {
return None;
}
if message.author.bot && !self.config.include_bot_messages {
return None;
}
let guild_id = message.guild_id.map(|id| id.to_string());
let channel_id = message.channel_id.to_string();
let channel_filter_enabled = !self.config.channel_ids.is_empty();
if channel_filter_enabled && !self.config.channel_ids.contains(&channel_id) {
return None;
}
self.note_seen_message(&channel_id, &message.id.to_string());
let mentions_bot = message
.mentions
.iter()
.any(|mention| mention.id == self.bot_user_id);
let is_dm = message.guild_id.is_none();
if !should_ingest_discord_message(
self.config.event_mode,
channel_filter_enabled,
is_dm,
mentions_bot,
) {
return None;
}
let conversation_id = self.conversation_id_for_channel(&channel_id);
let author_name = message
.author
.global_name
.as_deref()
.unwrap_or(&message.author.name)
.to_string();
let mut final_content = message.content.clone();
if !message.attachments.is_empty() {
let client = reqwest::Client::builder()
.user_agent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36")
.build()
.unwrap_or_else(|_| reqwest::Client::new());
for attachment in &message.attachments {
if attachment.size > 5 * 1024 * 1024 {
tracing::warn!(
filename = %attachment.filename,
size = attachment.size,
"Skipping Discord attachment: size exceeds 5MB limit"
);
continue;
}
tracing::info!(
filename = %attachment.filename,
url = %attachment.url,
"Downloading Discord attachment"
);
match client.get(&attachment.url).send().await {
Ok(resp) => {
if resp.status().is_success() {
match resp.bytes().await {
Ok(bytes) => {
use base64::{engine::general_purpose, Engine as _};
let b64 = general_purpose::STANDARD.encode(&bytes);
let extension = std::path::Path::new(&attachment.filename)
.extension()
.and_then(|ext| ext.to_str())
.map(|ext| ext.to_ascii_lowercase());
let mime = match extension.as_deref() {
Some("png") => "image/png".to_string(),
Some("jpg") | Some("jpeg") => "image/jpeg".to_string(),
Some("gif") => "image/gif".to_string(),
Some("webp") => "image/webp".to_string(),
Some("pdf") => "application/pdf".to_string(),
_ => attachment.content_type.clone().unwrap_or_else(|| {
"application/octet-stream".to_string()
}),
};
if !final_content.is_empty() {
final_content.push(' ');
}
final_content
.push_str(&format!("data:{};base64,{}", mime, b64));
}
Err(e) => {
tracing::warn!(?e, "Failed to read attachment bytes");
}
}
} else {
tracing::warn!(status = ?resp.status(), "Failed to download attachment");
}
}
Err(e) => {
tracing::warn!(?e, "Failed to fetch attachment");
}
}
}
}
Some(DiscordIncomingMessage {
conversation_id,
content: final_content,
author_id: message.author.id.to_string(),
author_name,
message_id: message.id.to_string(),
timestamp: message.timestamp.as_secs(),
guild_id,
channel_id,
mentions_bot,
is_dm,
author_is_bot: message.author.bot,
})
}
pub async fn send(&self, args: Value) -> Result<String> {
let channel_id: Id<ChannelMarker> = self.resolve_send_channel(&args)?;
let conversation_id = self.conversation_id_for_channel(&channel_id.to_string());
let content = string_arg(&args, "content")?;
validate_message_content(content)?;
let allowed_mentions = AllowedMentions::default();
let reference_message_id = self
.choose_reply_reference(&channel_id.to_string())
.map(|id| id.parse::<u64>())
.transpose()
.context("stored Discord reply target was not a snowflake")?
.map(Id::<MessageMarker>::new);
let mut request = self
.http
.create_message(channel_id)
.content(content)
.allowed_mentions(Some(&allowed_mentions));
if let Some(message_id) = reference_message_id {
request = request.reply(message_id);
}
let message = request
.await
.context("discord create message failed")?
.model()
.await
.context("failed to decode created Discord message")?;
Ok(format!(
"sent conv:{} msg:{} reply_to:{} time:{}",
conversation_id,
message.id,
reference_message_id
.map(|id| id.to_string())
.unwrap_or_else(|| "none".to_string()),
format_discord_timestamp(message.timestamp)
))
}
pub async fn fetch_history(&self, args: Value) -> Result<String> {
let channel_id: Id<ChannelMarker> = self.resolve_channel(&args, true)?;
let limit = args["limit"].as_u64().unwrap_or(20).clamp(1, 100) as u16;
let before_message_id =
optional_normalized_id_arg::<MessageMarker>(&args, "before_message_id")?;
let after_message_id =
optional_normalized_id_arg::<MessageMarker>(&args, "after_message_id")?;
if before_message_id.is_some() && after_message_id.is_some() {
anyhow::bail!("use only one of before_message_id or after_message_id");
}
let messages = match (before_message_id, after_message_id) {
(Some(id), None) => {
self.http
.channel_messages(channel_id)
.before(id)
.limit(limit)
.await?
.model()
.await?
}
(None, Some(id)) => {
self.http
.channel_messages(channel_id)
.after(id)
.limit(limit)
.await?
.model()
.await?
}
(None, None) => {
self.http
.channel_messages(channel_id)
.limit(limit)
.await?
.model()
.await?
}
(Some(_), Some(_)) => unreachable!(),
};
let conversation_id = self.existing_conversation_id_for_channel(&channel_id.to_string());
Ok(format_history(
channel_id,
conversation_id.as_deref(),
&messages,
))
}
pub async fn react(&self, args: Value) -> Result<String> {
let channel_id: Id<ChannelMarker> = self.resolve_channel(&args, true)?;
let message_id: Id<MessageMarker> = normalized_id_arg(&args, "message_id")?;
let emoji = string_arg(&args, "emoji")?;
let reaction = parse_reaction(emoji)?;
self.http
.create_reaction(channel_id, message_id, &reaction)
.await
.context("discord create reaction failed")?;
let conversation_id = self
.existing_conversation_id_for_channel(&channel_id.to_string())
.unwrap_or_else(|| channel_id.to_string());
Ok(format!("reacted conv:{conversation_id} msg:{message_id}"))
}
pub fn conversations(&self) -> Vec<(String, String, bool)> {
let state = self.state.lock().expect("discord runtime state poisoned");
let mut known = state
.conversation_by_channel
.iter()
.map(|(channel, conversation)| {
(
conversation.clone(),
channel.clone(),
state.active_channel_id.as_deref() == Some(channel.as_str()),
)
})
.collect::<Vec<_>>();
known.sort();
known
}
}
fn env_bool(name: &str) -> Option<bool> {
std::env::var(name).ok().map(|value| {
matches!(
value.to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
})
}
fn env_duration_secs(name: &str, default_secs: u64) -> Result<Duration> {
let Some(value) = std::env::var(name).ok() else {
return Ok(Duration::from_secs(default_secs));
};
let secs = value
.parse::<u64>()
.with_context(|| format!("{name} must be an integer number of seconds"))?;
Ok(Duration::from_secs(secs))
}
fn env_channel_ids(name: &str) -> Result<HashSet<String>> {
let Some(value) = std::env::var(name).ok() else {
return Ok(HashSet::new());
};
value
.split(|ch: char| ch == ',' || ch.is_ascii_whitespace())
.filter(|part| !part.is_empty())
.map(|part| {
part.parse::<u64>()
.with_context(|| format!("{name} contains invalid Discord channel id {part:?}"))?;
Ok(part.to_string())
})
.collect()
}
fn choose_reply_reference_from_state(state: &mut DiscordState, channel_id: &str) -> Option<String> {
let target = state.reply_target_by_channel.remove(channel_id)?;
let latest = state.latest_message_by_channel.get(channel_id);
if latest.is_some_and(|latest| latest != &target) {
Some(target)
} else {
None
}
}
fn conversation_id_for_channel_from_state(state: &mut DiscordState, channel_id: &str) -> String {
if let Some(id) = state.conversation_by_channel.get(channel_id) {
return id.clone();
}
state.next_conversation_id += 1;
let id = format!("d{}", state.next_conversation_id);
state
.conversation_by_channel
.insert(channel_id.to_string(), id.clone());
state
.channel_by_conversation
.insert(id.clone(), channel_id.to_string());
id
}
fn note_batch_targets_in_state<'a>(
state: &mut DiscordState,
latest_by_channel: HashMap<&'a str, &'a str>,
) {
state.active_channel_id = if latest_by_channel.len() == 1 {
latest_by_channel
.keys()
.next()
.map(|channel_id| (*channel_id).to_string())
} else {
None
};
for (channel_id, message_id) in latest_by_channel {
state
.reply_target_by_channel
.insert(channel_id.to_string(), message_id.to_string());
}
}
fn should_ingest_discord_message(
event_mode: DiscordEventMode,
channel_filter_enabled: bool,
is_dm: bool,
mentions_bot: bool,
) -> bool {
if is_dm || mentions_bot {
return true;
}
if channel_filter_enabled {
return true;
}
event_mode == DiscordEventMode::All
}
fn normalize_item_id(raw: &str) -> &str {
let s = raw.trim();
s.strip_prefix("msg:")
.or_else(|| s.strip_prefix("item:"))
.unwrap_or(s)
}
fn resolve_channel_arg_from_state(
state: &DiscordState,
args: &Value,
allow_active: bool,
) -> Result<String> {
if let Some(conversation_id) = optional_string_arg(args, "conversation_id")? {
return resolve_conversation_channel_from_state(state, conversation_id);
}
if let Some(channel_id) = optional_string_arg(args, "channel_id")? {
if parse_discord_id::<ChannelMarker>(channel_id, "Discord channel").is_ok() {
return Ok(channel_id.to_string());
}
return resolve_conversation_channel_from_state(state, channel_id).with_context(|| {
format!(
"channel_id {channel_id:?} is neither a Discord snowflake nor a known conversation id"
)
});
}
if allow_active {
if let Some(channel_id) = &state.active_channel_id {
return Ok(channel_id.clone());
}
}
anyhow::bail!(
"conversation_id or channel_id is required when the current Discord batch spans multiple conversations"
)
}
fn resolve_conversation_channel_from_state(
state: &DiscordState,
raw_conversation_id: &str,
) -> Result<String> {
let conversation_id = normalize_conversation_id(raw_conversation_id);
let Some(channel_id) = state.channel_by_conversation.get(conversation_id) else {
anyhow::bail!("unknown Discord conversation_id {conversation_id:?}");
};
Ok(channel_id.clone())
}
fn normalize_conversation_id(raw: &str) -> &str {
raw.strip_prefix("conv:").unwrap_or(raw)
}
fn normalized_id_arg<T>(args: &Value, key: &str) -> Result<Id<T>> {
let raw = string_arg(args, key)?;
parse_discord_id(normalize_item_id(raw), key)
}
fn optional_normalized_id_arg<T>(args: &Value, key: &str) -> Result<Option<Id<T>>> {
if args.get(key).is_none_or(Value::is_null) {
return Ok(None);
}
normalized_id_arg(args, key).map(Some)
}
fn parse_discord_id<T>(raw: &str, label: &str) -> Result<Id<T>> {
let id = raw
.parse::<u64>()
.with_context(|| format!("{label} must be a Discord snowflake"))?;
Ok(Id::new(id))
}
fn optional_string_arg<'a>(args: &'a Value, key: &str) -> Result<Option<&'a str>> {
if args.get(key).is_none_or(Value::is_null) {
return Ok(None);
}
string_arg(args, key).map(Some)
}
fn string_arg<'a>(args: &'a Value, key: &str) -> Result<&'a str> {
args[key]
.as_str()
.filter(|value| !value.trim().is_empty())
.ok_or_else(|| anyhow!("missing required arg '{key}'"))
}
fn validate_message_content(content: &str) -> Result<()> {
if content.chars().count() > 2000 {
anyhow::bail!("Discord message content exceeds 2000 characters");
}
Ok(())
}
fn parse_reaction(input: &str) -> Result<RequestReactionType<'_>> {
let emoji = input.trim();
if let Some(custom) = emoji.strip_prefix('<').and_then(|s| s.strip_suffix('>')) {
let parts = custom.split(':').collect::<Vec<_>>();
let (name, id) = match parts.as_slice() {
["", name, id] | ["a", name, id] => (*name, *id),
_ => anyhow::bail!("invalid custom emoji; expected <:name:id> or <a:name:id>"),
};
let id = id.parse::<u64>().with_context(|| {
let fallback = emoji_from_shortcodeish(name)
.map(|emoji| format!("; use unicode {emoji} instead if you just want :{name}:"))
.unwrap_or_default();
format!("custom emoji id must be a Discord snowflake{fallback}")
})?;
if id < 100_000_000_000_000_000 {
if let Some(emoji) = emoji_from_shortcodeish(name) {
return Ok(RequestReactionType::Unicode { name: emoji });
}
anyhow::bail!(
"custom emoji id is too short to look like a real Discord emoji id; use unicode emoji or copy an exact custom emoji string from Discord"
);
}
return Ok(RequestReactionType::Custom {
id: Id::<EmojiMarker>::new(id),
name: Some(name),
});
}
if let Some(emoji) = emoji_from_unicode_or_shortcode(emoji) {
return Ok(RequestReactionType::Unicode { name: emoji });
}
Ok(RequestReactionType::Unicode { name: emoji })
}
fn emoji_from_unicode_or_shortcode(input: &str) -> Option<&'static str> {
emojis::get(input)
.map(emojis::Emoji::as_str)
.or_else(|| emoji_from_shortcodeish(input))
}
fn emoji_from_shortcodeish(input: &str) -> Option<&'static str> {
let shortcode = normalize_emoji_shortcode(input)?;
emojis::get_by_shortcode(&shortcode).map(emojis::Emoji::as_str)
}
fn normalize_emoji_shortcode(input: &str) -> Option<String> {
let name = input.trim().trim_matches(':');
if name.is_empty()
|| name.chars().any(char::is_whitespace)
|| name
.chars()
.any(|ch| !(ch.is_ascii_alphanumeric() || matches!(ch, '_' | '-' | '+')))
{
return None;
}
Some(name.replace('-', "_").to_ascii_lowercase())
}