use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::net::TcpStream;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use tungstenite::stream::MaybeTlsStream;
use tungstenite::{connect, Message as WsMessage, WebSocket};
type GatewaySocket = WebSocket<MaybeTlsStream<TcpStream>>;
const API_BASE: &str = "https://discord.com/api/v10";
const GATEWAY_URL: &str = "wss://gateway.discord.gg/?v=10&encoding=json";
pub const MESSAGE_MAX: usize = 1900;
const DISPATCH_QUEUE_CAPACITY: usize = 32;
const INTENTS: u64 = (1 << 0) | (1 << 9) | (1 << 12) | (1 << 15);
const OP_DISPATCH: u64 = 0;
const OP_HEARTBEAT: u64 = 1;
const OP_IDENTIFY: u64 = 2;
const OP_RECONNECT: u64 = 7;
const OP_INVALID_SESSION: u64 = 9;
const OP_HELLO: u64 = 10;
const OP_HEARTBEAT_ACK: u64 = 11;
const CONFIG_RELATIVE_PATH: &str = ".yana-ai/os/discord-config.json";
const BOT_TOKEN_ENV_VAR: &str = "DISCORD_BOT_TOKEN";
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DiscordConfig {
#[serde(default)]
pub allowed_channel_ids: Vec<String>,
#[serde(default)]
pub allowed_user_ids: Vec<String>,
}
fn config_path(root: &Path) -> PathBuf {
root.join(CONFIG_RELATIVE_PATH)
}
pub fn load_config(root: &Path) -> Result<DiscordConfig> {
let path = config_path(root);
match std::fs::read_to_string(&path) {
Ok(text) => serde_json::from_str(&text)
.with_context(|| format!("invalid discord config at {}", path.display())),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(DiscordConfig::default()),
Err(error) => Err(error).with_context(|| format!("reading {}", path.display())),
}
}
pub fn bot_token() -> Option<String> {
std::env::var(BOT_TOKEN_ENV_VAR)
.ok()
.filter(|s| !s.trim().is_empty())
}
pub fn bot_token_env_var_name() -> &'static str {
BOT_TOKEN_ENV_VAR
}
pub fn is_allowed(cfg: &DiscordConfig, channel: &str, user: &str) -> bool {
cfg.allowed_channel_ids.iter().any(|id| id == channel)
&& (cfg.allowed_user_ids.is_empty() || cfg.allowed_user_ids.iter().any(|id| id == user))
}
pub struct Client {
agent: ureq::Agent,
token: String,
}
impl Client {
pub fn new(token: String) -> Self {
let config = ureq::Agent::config_builder()
.timeout_connect(Some(Duration::from_secs(10)))
.timeout_recv_response(Some(Duration::from_secs(15)))
.http_status_as_error(false)
.build();
Self {
agent: ureq::Agent::new_with_config(config),
token,
}
}
fn auth(&self) -> String {
format!("Bot {}", self.token)
}
pub fn send_message(&self, channel_id: &str, content: &str) -> Result<()> {
let url = format!("{API_BASE}/channels/{channel_id}/messages");
let mut resp = self
.agent
.post(&url)
.header("Authorization", self.auth())
.header("content-type", "application/json")
.send_json(json!({ "content": content }))
.context("discord sendMessage")?;
if resp.status().as_u16() >= 300 {
let body: String = resp
.body_mut()
.read_to_string()
.unwrap_or_default()
.chars()
.take(200)
.collect();
bail!(
"discord send failed: HTTP {} {}",
resp.status(),
body.trim()
);
}
Ok(())
}
pub fn get_me(&self) -> Result<String> {
let url = format!("{API_BASE}/users/@me");
let mut resp = self
.agent
.get(&url)
.header("Authorization", self.auth())
.call()
.context("discord getMe")?;
if resp.status().as_u16() >= 300 {
bail!("discord rejected the token (HTTP {})", resp.status());
}
let v: Value = resp.body_mut().read_json().context("parsing /users/@me")?;
Ok(v.get("username")
.and_then(Value::as_str)
.unwrap_or("?")
.to_string())
}
}
pub struct Incoming {
pub message_id: String,
pub channel_id: String,
pub user_id: String,
pub content: String,
}
fn ensure_crypto_provider_installed() {
static INSTALL: std::sync::Once = std::sync::Once::new();
INSTALL.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
pub fn run_gateway(
token: &str,
cfg: &DiscordConfig,
on_message: impl FnMut(Incoming) + Send + 'static,
) {
ensure_crypto_provider_installed();
let (tx, rx) = std::sync::mpsc::sync_channel::<Incoming>(DISPATCH_QUEUE_CAPACITY);
std::thread::spawn(move || {
let mut on_message = on_message;
while let Ok(incoming) = rx.recv() {
let result =
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| on_message(incoming)));
if let Err(panic) = result {
let message = panic
.downcast_ref::<&str>()
.map(|s| s.to_string())
.or_else(|| panic.downcast_ref::<String>().cloned())
.unwrap_or_else(|| "(non-string panic payload)".to_string());
eprintln!(
"[discord worker] a turn panicked and was recovered, the worker keeps \
running: {message}"
);
}
}
});
let mut backoff = Duration::from_secs(1);
loop {
let start = Instant::now();
match gateway_once(token, cfg, &tx) {
Ok(()) => backoff = Duration::from_secs(1),
Err(error) => {
if error.downcast_ref::<FatalCloseError>().is_some() {
eprintln!(
"[discord gateway] permanent failure: {error:#} — not reconnecting \
(fix the token / enable the required intent, then restart)."
);
return;
}
eprintln!("[discord gateway] {error:#}");
if start.elapsed() >= Duration::from_secs(60) {
backoff = Duration::from_secs(1);
}
std::thread::sleep(backoff);
backoff = (backoff * 2).min(Duration::from_secs(60));
}
}
}
}
fn gateway_once(
token: &str,
cfg: &DiscordConfig,
tx: &std::sync::mpsc::SyncSender<Incoming>,
) -> Result<()> {
let (mut ws, _) = connect(GATEWAY_URL).context("connecting to discord gateway")?;
set_read_timeout(&ws, Some(Duration::from_secs(20)));
let mut interval = Duration::from_millis(41_250);
let hello_deadline = Instant::now() + Duration::from_secs(20);
loop {
if Instant::now() >= hello_deadline {
bail!("timed out waiting for gateway HELLO (20s)");
}
let Some(v) = next_json(&mut ws)? else {
continue;
};
if v.get("op").and_then(Value::as_u64) == Some(OP_HELLO) {
if let Some(ms) = v
.get("d")
.and_then(|d| d.get("heartbeat_interval"))
.and_then(Value::as_u64)
{
interval = Duration::from_millis(ms.max(1000));
}
break;
}
}
ws.send(WsMessage::text(
json!({
"op": OP_IDENTIFY,
"d": {
"token": token,
"intents": INTENTS,
"properties": {"os": std::env::consts::OS, "browser": "yana-rt", "device": "yana-rt"}
}
})
.to_string(),
))
.context("sending IDENTIFY")?;
set_read_timeout(&ws, Some(Duration::from_millis(500)));
let mut last_seq: Option<u64> = None;
let mut awaiting_ack = false;
let mut next_beat = Instant::now() + interval;
loop {
if Instant::now() >= next_beat {
if awaiting_ack {
bail!("no Heartbeat ACK since the last beat — zombied gateway link, reconnecting");
}
ws.send(WsMessage::text(
json!({"op": OP_HEARTBEAT, "d": last_seq}).to_string(),
))
.context("sending heartbeat")?;
awaiting_ack = true;
next_beat = Instant::now() + interval;
}
let v = match next_json(&mut ws) {
Ok(Some(v)) => v,
Ok(None) => continue,
Err(error) => return Err(error),
};
if let Some(s) = v.get("s").and_then(Value::as_u64) {
last_seq = Some(s);
}
match v.get("op").and_then(Value::as_u64) {
Some(OP_HEARTBEAT) => {
ws.send(WsMessage::text(
json!({"op": OP_HEARTBEAT, "d": last_seq}).to_string(),
))
.context("sending requested heartbeat")?;
awaiting_ack = true;
}
Some(OP_HEARTBEAT_ACK) => awaiting_ack = false,
Some(OP_RECONNECT) => bail!("gateway requested reconnect (op 7)"),
Some(OP_INVALID_SESSION) => bail!("gateway invalidated the session (op 9)"),
Some(OP_DISPATCH) if v.get("t").and_then(Value::as_str) == Some("MESSAGE_CREATE") => {
if let Some(inc) = parse_message_create(v.get("d")) {
if is_allowed(cfg, &inc.channel_id, &inc.user_id) {
if let Err(std::sync::mpsc::TrySendError::Full(dropped)) = tx.try_send(inc)
{
eprintln!(
"[discord gateway] dispatch queue full ({DISPATCH_QUEUE_CAPACITY} \
pending) — dropping message from channel {} (worker is behind, \
likely mid-turn on a slow provider)",
dropped.channel_id
);
}
}
}
}
_ => {}
}
}
}
fn next_json(ws: &mut GatewaySocket) -> Result<Option<Value>> {
match ws.read() {
Ok(WsMessage::Text(text)) => Ok(serde_json::from_str(text.as_str()).ok()),
Ok(WsMessage::Ping(payload)) => {
let _ = ws.send(WsMessage::Pong(payload));
Ok(None)
}
Ok(WsMessage::Close(frame)) => {
let code = frame.as_ref().map(|f| u16::from(f.code));
if matches!(code, Some(4004 | 4010 | 4011 | 4012 | 4013 | 4014)) {
bail!(FatalCloseError(format!(
"gateway closed with permanent code {code:?} — bad token or a required \
intent (e.g. MESSAGE_CONTENT) is not enabled in the Developer Portal"
)));
}
bail!("gateway sent Close (code {code:?})");
}
Ok(_) => Ok(None),
Err(tungstenite::Error::Io(io_error))
if matches!(
io_error.kind(),
std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
) =>
{
Ok(None)
}
Err(error) => Err(error).context("gateway read"),
}
}
struct FatalCloseError(String);
impl std::fmt::Debug for FatalCloseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl std::fmt::Display for FatalCloseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl std::error::Error for FatalCloseError {}
fn set_read_timeout(ws: &GatewaySocket, timeout: Option<Duration>) {
let result = match ws.get_ref() {
MaybeTlsStream::Plain(stream) => stream.set_read_timeout(timeout),
MaybeTlsStream::Rustls(stream) => stream.sock.set_read_timeout(timeout),
_ => Ok(()),
};
let _ = result;
}
fn parse_message_create(d: Option<&Value>) -> Option<Incoming> {
let d = d?;
if d.get("author")
.and_then(|a| a.get("bot"))
.and_then(Value::as_bool)
.unwrap_or(false)
{
return None;
}
let message_id = d.get("id")?.as_str()?.to_string();
let channel_id = d.get("channel_id")?.as_str()?.to_string();
let user_id = d.get("author")?.get("id")?.as_str()?.to_string();
let content = d.get("content")?.as_str()?.trim().to_string();
if content.is_empty() {
return None;
}
Some(Incoming {
message_id,
channel_id,
user_id,
content,
})
}
pub fn chunk_reply(s: &str, max: usize) -> Vec<String> {
if s.encode_utf16().count() <= max {
return vec![s.to_string()];
}
let mut out = Vec::new();
let mut cur = String::new();
let mut units = 0usize;
for ch in s.chars() {
let u = ch.len_utf16();
if units + u > max && !cur.is_empty() {
out.push(std::mem::take(&mut cur));
units = 0;
}
cur.push(ch);
units += u;
}
if !cur.is_empty() {
out.push(cur);
}
out
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Condvar, Mutex};
#[test]
fn dispatch_never_blocks_while_the_worker_is_mid_turn() {
let (tx, rx) = std::sync::mpsc::channel::<Incoming>();
let processing_started = Arc::new((Mutex::new(false), Condvar::new()));
let ps = processing_started.clone();
std::thread::spawn(move || {
while let Ok(_incoming) = rx.recv() {
{
let (lock, cvar) = &*ps;
*lock.lock().unwrap() = true;
cvar.notify_all();
}
std::thread::sleep(Duration::from_millis(300));
}
});
let dummy = |id: &str| Incoming {
message_id: id.to_string(),
channel_id: "1".into(),
user_id: "1".into(),
content: "hi".into(),
};
tx.send(dummy("1")).unwrap();
let (lock, cvar) = &*processing_started;
let (_guard, timed_out) = cvar
.wait_timeout_while(lock.lock().unwrap(), Duration::from_secs(2), |started| {
!*started
})
.unwrap();
assert!(!timed_out.timed_out(), "worker never started processing");
let second_send_started = Instant::now();
tx.send(dummy("2")).unwrap();
let elapsed = second_send_started.elapsed();
assert!(
elapsed < Duration::from_millis(100),
"dispatching a new message must not block on a slow in-progress turn, took {elapsed:?}"
);
}
#[test]
fn worker_survives_a_panicking_turn_and_keeps_processing() {
let (tx, rx) = std::sync::mpsc::channel::<Incoming>();
let processed = Arc::new(Mutex::new(Vec::<String>::new()));
let processed_worker = processed.clone();
let handle = std::thread::spawn(move || {
while let Ok(incoming) = rx.recv() {
let id = incoming.message_id.clone();
let processed = processed_worker.clone();
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
if incoming.content == "boom" {
panic!("simulated turn failure");
}
processed.lock().unwrap().push(id.clone());
}));
if result.is_err() {
}
}
});
tx.send(Incoming {
message_id: "1".into(),
channel_id: "1".into(),
user_id: "1".into(),
content: "hello".into(),
})
.unwrap();
tx.send(Incoming {
message_id: "2".into(),
channel_id: "1".into(),
user_id: "1".into(),
content: "boom".into(),
})
.unwrap();
tx.send(Incoming {
message_id: "3".into(),
channel_id: "1".into(),
user_id: "1".into(),
content: "hello again".into(),
})
.unwrap();
drop(tx);
handle
.join()
.expect("the worker thread itself must not panic even though a turn inside it did");
assert_eq!(
*processed.lock().unwrap(),
vec!["1".to_string(), "3".to_string()],
"message 2 panicked and is correctly absent, but 1 and 3 must both \
still have been processed -- the panic must not have killed the loop"
);
}
#[test]
fn intents_include_message_content_and_guild_messages() {
assert_eq!(
INTENTS & (1 << 15),
1 << 15,
"MESSAGE_CONTENT must be requested"
);
assert_eq!(
INTENTS & (1 << 9),
1 << 9,
"GUILD_MESSAGES must be requested"
);
}
#[test]
fn gateway_opcodes_match_discord_v10() {
assert_eq!((OP_DISPATCH, OP_HEARTBEAT, OP_IDENTIFY), (0, 1, 2));
assert_eq!(
(OP_RECONNECT, OP_INVALID_SESSION, OP_HELLO, OP_HEARTBEAT_ACK),
(7, 9, 10, 11)
);
}
#[test]
fn allowlist_denies_empty_and_unlisted() {
let mut cfg = DiscordConfig::default();
assert!(
!is_allowed(&cfg, "100", "7"),
"empty channel list denies everyone"
);
cfg.allowed_channel_ids = vec!["100".into(), "200".into()];
assert!(
is_allowed(&cfg, "100", "7"),
"listed channel, no user restriction"
);
assert!(!is_allowed(&cfg, "300", "7"), "unlisted channel denied");
cfg.allowed_user_ids = vec!["7".into()];
assert!(is_allowed(&cfg, "100", "7"), "listed channel + listed user");
assert!(
!is_allowed(&cfg, "100", "8"),
"listed channel, unlisted user denied"
);
}
#[test]
fn parse_skips_bots_and_empty_parses_ids() {
let bot = json!({"id":"9","channel_id":"1","author":{"id":"2","bot":true},"content":"hi"});
assert!(
parse_message_create(Some(&bot)).is_none(),
"bot author skipped"
);
let empty = json!({"id":"9","channel_id":"1","author":{"id":"2"},"content":" "});
assert!(
parse_message_create(Some(&empty)).is_none(),
"empty content skipped"
);
let missing_id = json!({"channel_id":"1","author":{"id":"2"},"content":"hi"});
assert!(
parse_message_create(Some(&missing_id)).is_none(),
"a payload missing the message id must not parse"
);
let ok = json!({"id":"789","channel_id":"123","author":{"id":"456"},"content":"hello"});
let inc = parse_message_create(Some(&ok)).expect("valid message parses");
assert_eq!(
(
inc.message_id.as_str(),
inc.channel_id.as_str(),
inc.user_id.as_str(),
inc.content.as_str()
),
("789", "123", "456", "hello")
);
}
#[test]
fn chunk_reply_splits_over_the_limit_and_leaves_short_replies_whole() {
assert_eq!(chunk_reply("hello", 1900), vec!["hello".to_string()]);
let long = "a".repeat(10);
let chunks = chunk_reply(&long, 4);
assert_eq!(chunks, vec!["aaaa", "aaaa", "aa"]);
}
#[test]
fn load_config_parses_a_hand_edited_file() {
let root = std::env::temp_dir().join(format!("yana-discord-cfg-{}", uuid::Uuid::new_v4()));
std::fs::create_dir_all(root.join(".yana-ai/os")).unwrap();
std::fs::write(
config_path(&root),
r#"{"allowed_channel_ids": ["1", "2"], "allowed_user_ids": ["9"]}"#,
)
.unwrap();
let loaded = load_config(&root).unwrap();
assert_eq!(loaded.allowed_channel_ids, vec!["1", "2"]);
assert_eq!(loaded.allowed_user_ids, vec!["9"]);
std::fs::remove_dir_all(root).ok();
}
#[test]
fn load_config_defaults_to_empty_deny_all_when_no_file_exists() {
let root = std::env::temp_dir().join(format!("yana-discord-cfg-{}", uuid::Uuid::new_v4()));
let cfg = load_config(&root).unwrap();
assert!(cfg.allowed_channel_ids.is_empty());
assert!(
!is_allowed(&cfg, "1", "1"),
"no config file must still deny everyone"
);
}
#[test]
fn dispatch_queue_rejects_over_capacity_sends_without_blocking() {
let (tx, _rx) = std::sync::mpsc::sync_channel::<Incoming>(DISPATCH_QUEUE_CAPACITY);
let dummy = |id: &str| Incoming {
message_id: id.to_string(),
channel_id: "1".into(),
user_id: "1".into(),
content: "hi".into(),
};
let fill_started = Instant::now();
for i in 0..DISPATCH_QUEUE_CAPACITY {
tx.try_send(dummy(&i.to_string()))
.expect("queue must accept sends up to its own declared capacity");
}
assert!(
fill_started.elapsed() < Duration::from_millis(100),
"filling the queue to capacity must be fast — nothing is \
draining it, so any blocking here would hang the test"
);
match tx.try_send(dummy("overflow")) {
Err(std::sync::mpsc::TrySendError::Full(rejected)) => {
assert_eq!(rejected.message_id, "overflow");
}
Ok(()) => panic!("expected TrySendError::Full once the queue is at capacity, got Ok"),
Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
panic!("expected TrySendError::Full, got Disconnected — receiver was dropped?")
}
}
}
}