johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::time::Duration;

use async_trait::async_trait;
use serde_json::Value;
use tokio::sync::mpsc::UnboundedSender;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AtircIdentity {
    pub nick: String,
    pub handle: String,
    pub subhandle: String,
}

impl AtircIdentity {
    pub fn is_self(&self, identity: Option<&str>) -> bool {
        let Some(identity) = identity else {
            return false;
        };
        let identity = identity.to_ascii_lowercase();
        let identity = identity
            .split_once('@')
            .map_or(identity.as_str(), |(_, handle)| handle);
        if identity.eq_ignore_ascii_case(&self.subhandle) {
            return true;
        }
        self.subhandle.eq_ignore_ascii_case(&self.handle)
            && identity.eq_ignore_ascii_case(&self.handle)
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ConnectionInfo {
    pub subscribed_as: String,
    pub owner_nick: Option<String>,
    pub your_user_is: Option<String>,
    pub transport: Option<String>,
    pub identities: Vec<String>,
    pub channels: Vec<String>,
}

impl ConnectionInfo {
    pub fn from_value(value: &Value) -> Option<Self> {
        if value.get("event").and_then(Value::as_str) != Some("connected") {
            return None;
        }
        let subscribed_as = value.get("subscribedAs")?.as_str()?.to_string();
        let owner_nick = value
            .get("ownerNick")
            .and_then(Value::as_str)
            .map(str::to_string);
        let your_user_is = value
            .get("yourUserIs")
            .and_then(Value::as_str)
            .map(str::to_string);
        let transport = value
            .get("transport")
            .and_then(Value::as_str)
            .map(str::to_string);
        let identities = string_array(value.get("identities"));
        let channels = string_array(value.get("channels"));

        Some(Self {
            subscribed_as,
            owner_nick,
            your_user_is,
            transport,
            identities,
            channels,
        })
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ChannelMessage {
    pub channel: String,
    pub identity: Option<String>,
    pub text: String,
}

impl ChannelMessage {
    pub fn from_value(value: &Value) -> Option<Self> {
        if value.get("event").is_some() {
            return None;
        }
        let channel = value.get("channel")?.as_str()?.to_string();
        let text = value.get("text")?.as_str()?.to_string();
        Some(Self {
            channel,
            identity: value
                .get("identity")
                .and_then(Value::as_str)
                .map(str::to_string),
            text,
        })
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AtircEvent {
    Connected(ConnectionInfo),
    Message(ChannelMessage),
}

#[async_trait]
pub trait AtircClient: Send + Sync {
    fn identity(&self) -> AtircIdentity;
    async fn subscribe(
        &self,
        channels: Vec<String>,
        cursor: i64,
        tx: UnboundedSender<AtircEvent>,
    ) -> Result<(), String>;
    async fn publish(&self, channel: &str, text: &str) -> Result<(), String>;
}

pub struct LibraryAtircClient {
    endpoint: atirc_client::endpoint::Endpoint,
    token: String,
    identity: AtircIdentity,
}

impl LibraryAtircClient {
    pub fn from_token(token: impl Into<String>) -> Result<Self, String> {
        let token = token.into();
        let info = atirc_client::transport::token_info(&token)
            .map_err(|error| format!("invalid atirc token: {error:#}"))?;
        let endpoint = atirc_client::endpoint::Endpoint::parse(&info.endpoint)
            .map_err(|error| format!("invalid atirc endpoint: {error:#}"))?;
        Ok(Self {
            endpoint,
            token,
            identity: AtircIdentity {
                nick: info.nick,
                handle: info.handle,
                subhandle: info.subhandle,
            },
        })
    }
}

#[async_trait]
impl AtircClient for LibraryAtircClient {
    fn identity(&self) -> AtircIdentity {
        self.identity.clone()
    }

    async fn subscribe(
        &self,
        channels: Vec<String>,
        cursor: i64,
        tx: UnboundedSender<AtircEvent>,
    ) -> Result<(), String> {
        let (value_tx, mut value_rx) = tokio::sync::mpsc::unbounded_channel();
        let endpoint = self.endpoint.clone();
        let token = self.token.clone();
        let bridge = tokio::spawn(async move {
            if let Err(error) = atirc_client::transport::subscribe_messages(
                &endpoint, &token, &channels, cursor, value_tx,
            )
            .await
            {
                eprintln!(
                    "{}",
                    serde_json::json!({
                        "level": "error",
                        "event": "atirc.subscription_ended",
                        "error": format!("{error:#}"),
                    })
                );
            }
        });

        while let Some(value) = value_rx.recv().await {
            let event = if let Some(connection) = ConnectionInfo::from_value(&value) {
                Some(AtircEvent::Connected(connection))
            } else {
                ChannelMessage::from_value(&value).map(AtircEvent::Message)
            };
            if let Some(event) = event {
                if tx.send(event).is_err() {
                    break;
                }
            }
        }
        bridge.abort();
        Ok(())
    }

    async fn publish(&self, channel: &str, text: &str) -> Result<(), String> {
        for attempt in 0..3 {
            match atirc_client::transport::post(&self.endpoint, &self.token, Some(channel), text)
                .await
            {
                Ok(_) => return Ok(()),
                Err(error) if attempt < 2 => {
                    eprintln!(
                        "{}",
                        serde_json::json!({
                            "level": "warn",
                            "event": "atirc.publish_retry",
                            "attempt": attempt + 1,
                            "error": format!("{error:#}"),
                        })
                    );
                    tokio::time::sleep(Duration::from_millis(500)).await;
                }
                Err(error) => return Err(format!("{error:#}")),
            }
        }
        Err("atirc publish failed".to_string())
    }
}

pub fn synthetic_sub_command(channels: &[String], cursor: i64) -> String {
    let mut command = "atirc sub".to_string();
    for channel in channels {
        command.push_str(" -c ");
        command.push_str(&shell_quote(channel));
    }
    command.push_str(&format!(" --cursor {cursor}"));
    command
}

fn shell_quote(value: &str) -> String {
    if value
        .chars()
        .all(|character| character.is_ascii_alphanumeric() || "#_-./".contains(character))
    {
        value.to_string()
    } else {
        format!("'{}'", value.replace('\'', "'\\''"))
    }
}

fn string_array(value: Option<&Value>) -> Vec<String> {
    value
        .and_then(Value::as_array)
        .map(|items| {
            items
                .iter()
                .filter_map(Value::as_str)
                .map(str::to_string)
                .collect()
        })
        .unwrap_or_default()
}

fn is_nick_char(byte: u8) -> bool {
    byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'-'
}

pub fn mentions(text: &str, nick: &str) -> bool {
    if nick.is_empty() {
        return false;
    }
    let text = text.to_ascii_lowercase();
    let nick = nick.to_ascii_lowercase();
    let bytes = text.as_bytes();

    text.match_indices(&nick).any(|(start, _)| {
        let before_ok = start == 0 || !is_nick_char(bytes[start - 1]);
        let after = start + nick.len();
        let after_ok = after >= bytes.len() || !is_nick_char(bytes[after]);
        before_ok && after_ok
    })
}

/// Safe per-message payload. IRC delivers one PRIVMSG line capped at 512
/// bytes including the sender/channel prefix and CRLF; 400 bytes leaves room
/// for the longest normal prefix.
pub const ATIRC_MAX_MESSAGE_BYTES: usize = 400;

/// A reply that should publish nothing: empty, or any line that is just
/// `NO_REPLY` (case-insensitive; quotes, backticks, and periods ignored).
pub fn is_no_reply(text: &str) -> bool {
    text.trim().is_empty()
        || text.lines().any(|line| {
            line.trim()
                .trim_matches(|c: char| matches!(c, '`' | '"' | '\'' | '.' | '*'))
                .eq_ignore_ascii_case(crate::code_mode::watch::NO_REPLY)
        })
}

/// Split an assistant reply into IRC-sized messages.
///
/// Each non-empty source line becomes its own message. Lines longer than
/// [`ATIRC_MAX_MESSAGE_BYTES`] are wrapped at word boundaries; a word longer
/// than the limit is split at a UTF-8 boundary. Blank lines are dropped
/// because atirc rejects empty posts. A `NO_REPLY` reply publishes nothing.
pub fn format_reply(text: &str) -> Vec<String> {
    if is_no_reply(text) {
        return Vec::new();
    }
    let normalized = text.replace("\r\n", "\n").replace('\r', "\n");
    let mut messages = Vec::new();
    for line in normalized.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        wrap_line(line, ATIRC_MAX_MESSAGE_BYTES, &mut messages);
    }
    messages
}

fn wrap_line(line: &str, max_bytes: usize, messages: &mut Vec<String>) {
    let mut remaining = line;
    while remaining.len() > max_bytes {
        let mut cut = max_bytes;
        while cut > 0 && !remaining.is_char_boundary(cut) {
            cut -= 1;
        }
        let word_boundary = remaining[..cut]
            .rfind(char::is_whitespace)
            .filter(|index| *index > 0);
        let (head, tail) = match word_boundary {
            Some(index) => (&remaining[..index], &remaining[index..]),
            None => (&remaining[..cut], &remaining[cut..]),
        };
        let head = head.trim_end();
        if !head.is_empty() {
            messages.push(head.to_string());
        }
        remaining = tail.trim_start();
    }
    if !remaining.is_empty() {
        messages.push(remaining.to_string());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_connection_info() {
        let value = serde_json::json!({
            "event": "connected",
            "transport": "http",
            "subscribedAs": "alice",
            "ownerNick": "owner",
            "yourUserIs": "example.test",
            "channels": ["#test-asyncai"],
            "identities": ["alice@alice.example.test", "bob@bob.example.test"],
        });
        let info = ConnectionInfo::from_value(&value).unwrap();
        assert_eq!(info.subscribed_as, "alice");
        assert_eq!(info.owner_nick.as_deref(), Some("owner"));
        assert_eq!(info.transport.as_deref(), Some("http"));
        assert_eq!(info.channels, vec!["#test-asyncai"]);
        assert_eq!(
            info.identities,
            vec!["alice@alice.example.test", "bob@bob.example.test"]
        );
    }

    #[test]
    fn parses_message_and_ignores_events() {
        let value = serde_json::json!({
            "channel": "#test-asyncai",
            "identity": "alice.example.test",
            "ts": "2026-01-01T00:00:00Z",
            "text": "hello",
        });
        assert_eq!(
            ChannelMessage::from_value(&value),
            Some(ChannelMessage {
                channel: "#test-asyncai".to_string(),
                identity: Some("alice.example.test".to_string()),
                text: "hello".to_string(),
            })
        );
        assert_eq!(
            ChannelMessage::from_value(&serde_json::json!({"event": "join"})),
            None
        );
    }

    #[test]
    fn identity_self_filter_handles_handle_and_subhandle() {
        let identity = AtircIdentity {
            nick: "alice".to_string(),
            handle: "example.test".to_string(),
            subhandle: "alice.example.test".to_string(),
        };
        assert!(identity.is_self(Some("alice.example.test")));
        assert!(identity.is_self(Some("alice@alice.example.test")));
        assert!(!identity.is_self(Some("example.test")));
        assert!(!identity.is_self(Some("bob.example.test")));

        let account = AtircIdentity {
            nick: "john".to_string(),
            handle: "example.test".to_string(),
            subhandle: "example.test".to_string(),
        };
        assert!(account.is_self(Some("example.test")));
    }

    #[test]
    fn synthetic_command_describes_subscription() {
        assert_eq!(
            synthetic_sub_command(&["#test-asyncai".to_string()], -1),
            "atirc sub -c #test-asyncai --cursor -1"
        );
        assert_eq!(
            synthetic_sub_command(&["chan nel".to_string()], 5),
            "atirc sub -c 'chan nel' --cursor 5"
        );
    }

    #[test]
    fn mention_matching_uses_word_boundaries() {
        assert!(mentions("alice: tell bob", "alice"));
        assert!(mentions("@alice hello", "alice"));
        assert!(mentions("bob -> alice: hi", "alice"));
        assert!(!mentions("malice at the palace", "alice"));
        assert!(!mentions("bob: hello", "alice"));
    }

    #[test]
    fn reply_splits_multiline_text_into_irc_messages() {
        assert_eq!(
            format_reply("hello\nthere\n\n  spaced  \r\n"),
            vec![
                "hello".to_string(),
                "there".to_string(),
                "spaced".to_string()
            ]
        );
        assert!(format_reply("   \n\n").is_empty());
    }

    #[test]
    fn reply_wraps_long_lines_at_word_boundaries() {
        let line = "word ".repeat(200);
        let messages = format_reply(line.trim());
        assert!(messages.len() > 1);
        assert!(messages
            .iter()
            .all(|message| message.len() <= ATIRC_MAX_MESSAGE_BYTES));
        assert_eq!(messages.join(" "), line.trim());
    }

    #[test]
    fn reply_wraps_overlong_words_without_splitting_utf8() {
        let line = "é".repeat(ATIRC_MAX_MESSAGE_BYTES);
        let messages = format_reply(&line);
        assert!(messages.len() > 1);
        assert!(messages
            .iter()
            .all(|message| message.len() <= ATIRC_MAX_MESSAGE_BYTES));
        assert_eq!(messages.concat(), line);
    }
}