Fix Discord-to-Discord replies

Hyperfixation commit lol

mid-chat/src/lib.rs

@@ -13,11 +13,11 @@ pub use content::*;
 
 #[derive(Debug, Clone)]
 pub struct ChatMessage {
-    pub origin: ChatReference,
+    pub origin: ChatMessageReference,
     pub author: ChatAuthor,
     pub content: ChatMessageContent,
     pub attachments: Vec<()>,
-    pub replying: Option<ChatReference>,
+    pub replying: Option<ChatMessageReference>,
 }
 
 pub mod event;

mid-chat/src/reference.rs

@@ -3,3 +3,18 @@ pub struct ChatReference {
     pub service: &'static str,
     pub id: String,
 }
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ChatMessageReference {
+    pub channel: ChatReference,
+    pub message_id: String,
+}
+
+impl ChatMessageReference {
+    pub fn new(channel: ChatReference, message_id: impl ToString) -> Self {
+        Self {
+            channel,
+            message_id: message_id.to_string(),
+        }
+    }
+}

phoebe-main/src/main.rs

@@ -2,7 +2,49 @@ use color_eyre::Result;
 use tracing::info;
 use tracing_subscriber::EnvFilter;
 
-use phoebe::service::Service;
+use phoebe::{
+    get_linked_channels, link_messages,
+    prelude::{ChatEvent, SqlitePool},
+    service::Service,
+    DynServiceLookup,
+};
+
+async fn handle_events(
+    dyn_service: DynServiceLookup,
+    db: SqlitePool,
+    mut service: Box<dyn Service + Send + Sync>,
+    mut rx: tokio::sync::broadcast::Receiver<ChatEvent>,
+) {
+    info!("Handling events for {}…", service.tag());
+
+    let mut conn = db
+        .acquire()
+        .await
+        .expect("Failed to acquire core DB connection");
+
+    while let Ok(event) = rx.recv().await {
+        match event {
+            phoebe::prelude::ChatEvent::NewMessage(message) => {
+                let linked_channels =
+                    get_linked_channels(&mut conn, dyn_service, &message.origin.channel).await;
+
+                let mut resulting_messages = vec![];
+                for destination_channel in linked_channels {
+                    resulting_messages.extend(
+                        service
+                            .send_chat_message(&message, destination_channel)
+                            .await,
+                    )
+                }
+
+                if let Err(e) = link_messages(&mut conn, &message.origin, &resulting_messages).await
+                {
+                    tracing::error!("Failed to link messages: {e}");
+                }
+            }
+        }
+    }
+}
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -16,19 +58,22 @@ async fn main() -> Result<()> {
     let (tx, _) = tokio::sync::broadcast::channel(512);
 
     let db = phoebe::open_core_db().await?;
+
+    fn dyn_service(service: &str) -> &'static str {
+        match service {
+            "discord" => "discord",
+            "matrix" => "matrix",
+            _ => panic!("Unsupported service: {}", service),
+        }
+    }
+
     let services: Vec<Box<dyn Service + Send + Sync>> = vec![Box::new(
-        phoebe_discord::setup(db.clone(), tx.clone()).await?,
+        phoebe_discord::setup(db.clone(), tx.clone(), dyn_service).await?,
     )];
 
-    let handles = services.into_iter().map(|mut srv| {
-        let mut rx = tx.subscribe();
-        tokio::spawn(async move {
-            info!("Handling events for {}…", srv.get_service_tag());
-            while let Ok(event) = rx.recv().await {
-                srv.handle_chat_event(&event).await;
-            }
-        })
-    });
+    let handles = services
+        .into_iter()
+        .map(|srv| tokio::spawn(handle_events(dyn_service, db.clone(), srv, tx.subscribe())));
 
     let _ = futures::future::join_all(handles).await;
 

phoebe/.env

@@ -1 +1 @@
-DATABASE_URL="sqlite://PHOEBE_DB_ROOT/main.db"
+DATABASE_URL="sqlite://${PHOEBE_DB_ROOT}/main.db"

phoebe/Cargo.toml

@@ -11,3 +11,4 @@ sqlx = { version = "0.5", features = ["runtime-tokio-native-tls", "sqlite"] }
 tracing = "0.1"
 async-trait = "0.1.53"
 eyre = "0.6.8"
+futures = "0.3.21"

phoebe/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    println!("cargo:rerun-if-changed=migrations");
+}

phoebe/migrations/20220416061935_create_message_links_table.sql

@@ -0,0 +1,10 @@
+CREATE TABLE message_links (
+  id INTEGER PRIMARY KEY AUTOINCREMENT
+) STRICT;
+
+CREATE TABLE messages (
+  link_id INTEGER NOT NULL REFERENCES message_links(id), 
+  service TEXT NOT NULL,
+  channel TEXT NOT NULL,
+  message TEXT NOT NULL
+)

phoebe/migrations/20220416065641_create_channel_links_table.sql

@@ -0,0 +1,8 @@
+CREATE TABLE channel_links (
+  from_service TEXT NOT NULL,
+  from_channel TEXT NOT NULL,
+  to_service TEXT NOT NULL,
+  to_channel TEXT NOT NULL,
+
+  PRIMARY KEY (from_service, from_channel, to_service, to_channel)
+) STRICT;

phoebe/src/lib.rs

@@ -1,14 +1,135 @@
+use futures::{stream::BoxStream, Future};
 pub use mid_chat;
+use mid_chat::{ChatMessageReference, ChatReference};
+
+use futures::StreamExt;
+use sqlx::{Row, SqliteConnection, SqlitePool};
+use tokio::sync::broadcast::*;
 
 pub mod db;
 pub mod prelude;
 pub mod service;
 
-pub type ChatEventSender = tokio::sync::broadcast::Sender<mid_chat::event::ChatEvent>;
-pub type ChatEventReceiver = tokio::sync::broadcast::Receiver<mid_chat::event::ChatEvent>;
+pub type ChatEventSender = Sender<mid_chat::event::ChatEvent>;
+pub type ChatEventReceiver = Receiver<mid_chat::event::ChatEvent>;
 
-pub async fn open_core_db() -> sqlx::Result<sqlx::SqlitePool> {
+pub type DynServiceLookup = fn(&str) -> &'static str;
+
+pub async fn open_core_db() -> sqlx::Result<SqlitePool> {
     let db = db::open("main").await?;
     sqlx::migrate!().run(&db).await?;
     Ok(db)
 }
+
+pub async fn get_linked_channels(
+    conn: &mut SqliteConnection,
+    dyn_service: DynServiceLookup,
+    channel: &ChatReference,
+) -> Vec<ChatReference> {
+    let from_service = channel.service;
+    let from_channel = &channel.id;
+    let query = sqlx::query!(
+        "SELECT * FROM channel_links WHERE from_service = ? AND from_channel = ?",
+        from_service,
+        from_channel
+    );
+
+    query
+        .fetch(&mut *conn)
+        .filter_map(|r| async { r.ok() })
+        .map(|r| ChatReference {
+            service: dyn_service(&r.to_service),
+            id: r.to_channel,
+        })
+        .collect()
+        .await
+}
+
+pub async fn link_messages(
+    conn: &mut SqliteConnection,
+    origin: &ChatMessageReference,
+    messages: &[ChatMessageReference],
+) -> sqlx::Result<()> {
+    let message_link = sqlx::query!("INSERT INTO message_links DEFAULT VALUES")
+        .execute(&mut *conn)
+        .await?
+        .last_insert_rowid();
+
+    for resultant in std::iter::once(origin).chain(messages) {
+        let service = &resultant.channel.service;
+        let channel = &resultant.channel.id;
+        let message = &resultant.message_id;
+
+        let query = sqlx::query!(
+            "INSERT INTO messages VALUES (?, ?, ?, ?)",
+            message_link,
+            service,
+            channel,
+            message
+        );
+
+        let _ = query.execute(&mut *conn).await?;
+    }
+
+    Ok(())
+}
+
+pub async fn get_linked_messages<'a>(
+    conn: &'a mut SqliteConnection,
+    dyn_service: DynServiceLookup,
+    message: &ChatMessageReference,
+) -> sqlx::Result<BoxStream<'a, ChatMessageReference>> {
+    let link_id = {
+        let service = &message.channel.service;
+        let channel = &message.channel.id;
+        let message_id = &message.message_id;
+        let query = sqlx::query!(
+            "SELECT link_id FROM messages WHERE service = ? AND channel = ? AND message = ?",
+            service,
+            channel,
+            message_id
+        );
+        query.fetch_one(&mut *conn).await
+    }?
+    .link_id;
+
+    let stream = sqlx::query("SELECT * FROM messages WHERE link_id = ?")
+        .bind(link_id)
+        .fetch(&mut *conn)
+        .filter_map(|r| futures::future::ready(r.ok()))
+        .map(move |r| {
+            ChatMessageReference::new(
+                ChatReference {
+                    service: dyn_service(&r.get::<String, _>("service")),
+                    id: r.get("channel"),
+                },
+                r.get::<String, _>("message"),
+            )
+        });
+
+    Ok(Box::pin(stream))
+}
+
+pub async fn lookup_message<F, Fut>(
+    conn: &mut SqliteConnection,
+    dyn_service: DynServiceLookup,
+    linked_message: &ChatMessageReference,
+    filter: F,
+) -> Option<ChatMessageReference>
+where
+    F: FnMut(&ChatMessageReference) -> Fut,
+    Fut: Future<Output = bool>,
+{
+    let references = get_linked_messages(&mut *conn, dyn_service, linked_message)
+        .await
+        .ok()?
+        .filter(filter)
+        .collect::<Vec<_>>()
+        .await;
+
+    if let [reference] = references.as_slice() {
+        Some(reference.clone())
+    } else {
+        None
+    }
+}

phoebe/src/prelude.rs

@@ -1,6 +1,7 @@
 pub use crate::{service::Service, ChatEventReceiver, ChatEventSender};
 
 pub use async_trait::async_trait;
-pub use eyre::Result;
+pub use eyre::{self, Result};
+pub use futures::{self, prelude::*};
 pub use mid_chat::event::ChatEvent;
-pub use sqlx::SqlitePool;
+pub use sqlx::{self, SqliteConnection, SqlitePool};

phoebe/src/service.rs

@@ -1,7 +1,11 @@
-use mid_chat::event::ChatEvent;
+use mid_chat::{ChatMessage, ChatMessageReference, ChatReference};
 
 #[async_trait::async_trait]
 pub trait Service {
-    fn get_service_tag(&self) -> &'static str;
-    async fn handle_chat_event(&mut self, event: &ChatEvent);
+    fn tag(&self) -> &'static str;
+    async fn send_chat_message(
+        &mut self,
+        source: &ChatMessage,
+        destination_channel: ChatReference,
+    ) -> Vec<ChatMessageReference>;
 }

services/phoebe-discord/.env

@@ -1 +1 @@
-DATABASE_URL="sqlite://PHOEBE_DB_ROOT/discord_media.db"
+DATABASE_URL="sqlite://${PHOEBE_DB_ROOT}/discord_media.db"

services/phoebe-discord/src/chat_conv.rs

@@ -69,7 +69,7 @@ fn mid_to_discord(component: &ChatContentComponent) -> String {
                 lang.as_ref()
                     .map(|s| s.to_string())
                     .unwrap_or_else(|| "".to_string()),
-                source.to_string()
+                source
             )
         }
 

services/phoebe-discord/src/handler.rs

@@ -44,20 +44,17 @@ impl EventHandler for DiscordHandler {
     }
 
     async fn message(&self, ctx: Context, message: Message) {
-        let origin = ChatReference {
-            service: "discord",
-            id: message.id.to_string(),
-        };
+        let origin = ChatMessageReference::new(discord_reference(message.channel_id), message.id);
 
         let author = self.get_author(&ctx, &message).await;
 
         let content = discord_message_format::parse(&message.content);
         let content = super::chat_conv::convert(&content);
 
-        let replies_to = message.referenced_message.as_ref().map(|m| ChatReference {
-            service: "discord",
-            id: m.id.to_string(),
-        });
+        let replies_to = message
+            .referenced_message
+            .as_ref()
+            .map(|m| ChatMessageReference::new(discord_reference(m.channel_id), m.id));
 
         let chat_message = ChatMessage {
             origin,

services/phoebe-discord/src/lib.rs

@@ -1,9 +1,15 @@
-use phoebe::{mid_chat, prelude::*};
+use phoebe::{
+    mid_chat::{self, ChatMessage, ChatMessageReference, ChatReference},
+    prelude::*,
+    DynServiceLookup,
+};
 use serenity::{client::Context, Client};
 use tracing::{debug, info};
 
 mod chat_conv;
 mod handler;
+mod lookup;
+mod sender;
 
 pub fn discord_reference(id: impl ToString) -> mid_chat::ChatReference {
     mid_chat::ChatReference {
@@ -13,10 +19,17 @@ pub fn discord_reference(id: impl ToString) -> mid_chat::ChatReference {
 }
 
 pub struct DiscordService {
-    pub discord_ctx: Context,
+    pub core_db: SqlitePool,
+    pub discord_media_db: SqlitePool,
+    pub ctx: Context,
+    pub dyn_service: DynServiceLookup,
 }
 
-pub async fn setup(core_db: SqlitePool, tx: ChatEventSender) -> Result<DiscordService> {
+pub async fn setup(
+    core_db: SqlitePool,
+    tx: ChatEventSender,
+    dyn_service: DynServiceLookup,
+) -> Result<DiscordService> {
     info!("Setting up Discord service…");
 
     let discord_media_db = phoebe::db::open("discord_media").await?;
@@ -25,13 +38,12 @@ pub async fn setup(core_db: SqlitePool, tx: ChatEventSender) -> Result<DiscordSe
     let (ctx_tx, mut ctx_rx) = tokio::sync::mpsc::unbounded_channel::<Context>();
 
     let discord_handler = handler::DiscordHandler {
-        core_db,
-        discord_media_db,
+        core_db: core_db.clone(),
+        discord_media_db: discord_media_db.clone(),
         chat_event_tx: tx,
         ctx_tx,
     };
 
-    // TODO: Create a discord client
     debug!("Logging in…");
     let discord_token = std::env::var("PHOEBE_DISCORD_TOKEN")
         .expect("PHOEBE_DISCORD_TOKEN environment variable was not set!");
@@ -50,16 +62,30 @@ pub async fn setup(core_db: SqlitePool, tx: ChatEventSender) -> Result<DiscordSe
     let discord_ctx = ctx_rx.recv().await.expect("Couldn't get Discord context");
     debug!("Logged in!");
 
-    Ok(DiscordService { discord_ctx })
+    Ok(DiscordService {
+        core_db,
+        discord_media_db,
+        ctx: discord_ctx,
+        dyn_service,
+    })
 }
 
 #[async_trait]
 impl Service for DiscordService {
-    async fn handle_chat_event(&mut self, event: &ChatEvent) {
-        dbg!(event);
+    async fn send_chat_message(
+        &mut self,
+        source: &ChatMessage,
+        destination_channel: ChatReference,
+    ) -> Vec<ChatMessageReference> {
+        assert_eq!(destination_channel.service, "discord");
+        sender::send_discord_message(self, source, destination_channel)
+            .await
+            .ok()
+            .into_iter()
+            .collect()
     }
 
-    fn get_service_tag(&self) -> &'static str {
+    fn tag(&self) -> &'static str {
         "discord"
     }
 }

services/phoebe-discord/src/lookup.rs

@@ -0,0 +1,18 @@
+use phoebe::{lookup_message, mid_chat::ChatMessageReference, prelude::*};
+
+use crate::DiscordService;
+
+impl DiscordService {
+    pub async fn lookup_message<F, Fut>(
+        &self,
+        linked_message: &ChatMessageReference,
+        filter: F,
+    ) -> Option<ChatMessageReference>
+    where
+        F: FnMut(&ChatMessageReference) -> Fut,
+        Fut: Future<Output = bool>,
+    {
+        let mut conn = self.core_db.acquire().await.ok()?;
+        lookup_message(&mut conn, self.dyn_service, linked_message, filter).await
+    }
+}

services/phoebe-discord/src/sender.rs

@@ -0,0 +1,55 @@
+use phoebe::{
+    mid_chat::{ChatMessage, ChatMessageReference, ChatReference},
+    prelude::{future, Result},
+};
+
+use serenity::{model::prelude::*, prelude::*};
+
+use crate::{chat_conv, discord_reference, DiscordService};
+
+pub async fn send_discord_message(
+    discord: &mut DiscordService,
+    source: &ChatMessage,
+    destination_channel: ChatReference,
+) -> Result<ChatMessageReference> {
+    // TODO: If we have a webhook we should use the webhook
+
+    let channel_id = destination_channel.id.parse::<ChannelId>()?;
+    let formatted_message = chat_conv::format(&source.content);
+    let content = format!(
+        "{} ({}): {}",
+        source.author.display_name, source.author.reference.service, formatted_message
+    );
+
+    let discord_reply = if let Some(reply) = &source.replying {
+        if let Some(reply_ref) = discord
+            .lookup_message(reply, |r| future::ready(r.channel == destination_channel))
+            .await
+        {
+            assert_eq!(reply_ref.channel.service, "discord");
+            let channel_id: ChannelId = reply_ref.channel.id.parse().unwrap();
+            let message_id: MessageId = reply_ref.message_id.parse::<u64>().unwrap().into();
+            Some((channel_id, message_id))
+        } else {
+            None
+        }
+    } else {
+        None
+    };
+
+    let sent_message = channel_id
+        .send_message(&discord.ctx, move |m| {
+            let m = m.content(content);
+            if let Some(reply) = discord_reply {
+                m.reference_message(reply)
+            } else {
+                m
+            }
+        })
+        .await?;
+
+    Ok(ChatMessageReference::new(
+        discord_reference(sent_message.channel_id),
+        sent_message.id,
+    ))
+}