Build a custom messaging provider for Discord bots

A couple weekends ago I found myself wanting to write a Discord bot: specifically, a simple bot that I could tag with a message to trigger some processing and a response message. As I started thinking about what this bot could do, I realized this was an excellent use-case for a provider: instead of writing a one-off bot, I could create a reusable provider that would make it trivial to build any Discord bot I wanted.

The Discord API is fully-featured and massive, but for such a simple use-case, it would be overkill to work from scratch every time you wanted to write a bot. With a custom provider, I could simply write a component in any language and then plug it in to my provider via a simple API.

In this post, I'll retrace my steps to implement the custom provider using the wasmcloud-messaging interface, and show you how you can use the provider to write Discord bots of your own.

Getting started

wasmCloud supports two types of interfaces: well-known and custom. Well-known interfaces include everything in WASI 0.2 (like wasi-http and wasi-cli). It also includes interfaces from proposals (such as wasi-keyvalue and wasi-blobstore from the wasi-cloud proposal) and core wasmCloud functionality like wasmcloud-messaging for pubsub messaging.

We'll start off with the provider create page in the documentation, which already includes a template for a provider that implements wasmcloud-messaging. Clearing out the NATS-specific bits, we're left with a fairly barebones Rust program:

rust

use std::{collections::HashMap, sync::Arc};

use anyhow::Context as _;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;

use wasmcloud_provider_sdk::{get_connection, run_provider, Context, Provider};
use wit_bindgen_wrpc::tracing::{debug, error, warn};

wit_bindgen_wrpc::generate!();

use crate::wasmcloud::messaging::types;

#[derive(Default, Clone)]
struct DiscordProvider {}

impl Provider for DiscordProvider {
    async fn receive_link_config_as_source(
        &self,
        config: wasmcloud_provider_sdk::LinkConfig<'_>,
    ) -> Result<(), anyhow::Error> {
        todo!("implement receive link config")
    }

    async fn delete_link(&self, component_id: &str) -> Result<(), anyhow::Error> {
        todo!("implement delete link")
    }
}

impl exports::wasmcloud::messaging::consumer::Handler<Option<Context>> for DiscordProvider {
    async fn request(
        &self,
        _: Option<Context>,
        _: String,
        _: Vec<u8>,
        _: u32,
    ) -> Result<Result<types::BrokerMessage, String>, anyhow::Error> {
        todo!("implement request")
    }

    async fn publish(
        &self,
        ctx: Option<Context>,
        msg: types::BrokerMessage,
    ) -> Result<Result<(), String>, anyhow::Error> {
        todo!("implement publish")
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let provider = DiscordProvider::new();
    let shutdown = run_provider(provider.clone(), "messaging-discord-provider")
        .await
        .context("failed to run provider")?;
    let connection = get_connection();
    serve(
        &connection.get_wrpc_client(connection.provider_key()),
        provider,
        shutdown,
    )
    .await
}

Our loose plan for this provider is:

1. Whenever the provider is linked to a component to send messages using wasmcloud-messaging, look for an authentication token in the configuration and create a Discord client
2. Cache the client in-memory so that we can easily access it for modifications or shutdowns
3. Start a background task to forward any messages received by the client to the component, using its component ID
4. Implement publish as a mechanism to send messages on behalf of the bot

Taking a test-driven development approach, we can write a component using wasmcloud-messaging with my ideal workflow to show how simple it should be, and make it a success criterion for the provider to ensure the component worked. This bot should use the messaging contract to receive a message from Discord, see if the message included ping, and respond with a message back.

Oh, we'll also name our component after Janet from The Good Place, because why not?

rust

wit_bindgen::generate!();

use exports::wasmcloud::messaging::handler::Guest;
use wasmcloud::messaging::*;

struct GoodJanet;

// Yes, from the Good Place
impl Guest for GoodJanet {
    fn handle_message(msg: types::BrokerMessage) -> Result<(), String> {
        let content = String::from_utf8_lossy(&msg.body);
        if content.contains("ping") && msg.reply_to.is_some() {
            consumer::publish(&types::BrokerMessage {
                subject: msg.reply_to.unwrap(),
                reply_to: None,
                body: b"Ping received. What do you need help with?".to_vec(),
            })
        } else {
            Ok(())
        }
    }
}

export!(GoodJanet);

info

If we wish, we can test this bot directly with the NATS implementation of wasmcloud-messaging to first validate the functionality works as we expect. Just follow the testing the provider section of the provider developer guide with your component.

Implementation

Requirements

To implement a Discord bot, you will need to create an application on their developer portal and have a server you can install applications into.

When I started out originally, I decided to iterate quickly—as if this was just a barebones Rust binary, following the same process I would if I decided to write a Discord bot in Rust. My goal was to implement the functionality that I wanted, and then conform that functionality to the wasmCloud interface, so I can easily reuse the provider for each bot.

After a few searches on https://crates.io, I stumbled upon serenity which seemed like a good fit for my ideal use-case. I took the example bot directly from their documentation and put the logic into src/main.rs. We'll do the same thing now and comment out the provider code for the time being.

rust

#[group]
#[commands(ping)]
struct General;

#[tokio::main]
async fn main() {
    // TODO: Uncomment provider code
    // let provider = DiscordProvider::new();
    // let shutdown = run_provider(provider.clone(), "messaging-discord-provider")
    //     .await
    //     .context("failed to run provider")?;
    // let connection = get_connection();
    // serve(
    //     &connection.get_wrpc_client(connection.provider_key()),
    //     provider,
    //     shutdown,
    // )
    // .await
    let framework = StandardFramework::new().group(&GENERAL_GROUP);
    framework.configure(Configuration::new().prefix("~")); // set the bot's prefix to "~"

    // Login with a bot token from the environment
    let token = env::var("DISCORD_TOKEN").expect("token");
    let intents = GatewayIntents::non_privileged() | GatewayIntents::MESSAGE_CONTENT;
    let mut client = Client::builder(token, intents)
        .event_handler(Handler)
        .framework(framework)
        .await
        .expect("Error creating client");

    // start listening for events by starting a single shard
    if let Err(why) = client.start().await {
        println!("An error occurred while running the client: {:?}", why);
    }
}

#[command]
async fn ping(ctx: &Context, msg: &Message) -> CommandResult {
    msg.reply(ctx, "Pong!").await?;

    Ok(())
}

Using this, we can set the DISCORD_TOKEN environment variable and run the provider directly, sending a message to our provider and getting a "Pong!" in return. This gives us simple logic to start a client and handle messages. Taking a further look into the documentation, I learned that the framework concept wasn't required and instead I was able to simply supply an event handler, which will come in handy later.

Refactoring the existing code, we can move the client construction logic to receive_link_config_as_source and start to forward messages to the linked component when we receive a message from Discord:

rust

#[derive(Default, Clone)]
struct DiscordProvider {}

impl Provider for DiscordProvider {
    async fn receive_link_config_as_source(
        &self,
        config: wasmcloud_provider_sdk::LinkConfig<'_>,
    ) -> Result<(), anyhow::Error> {
        debug!("receiving link configuration as source");
        let source_config: case_insensitive_hashmap::CaseInsensitiveHashMap<String> =
            case_insensitive_hashmap::CaseInsensitiveHashMap::from_iter(
                config
                    .config
                    .iter()
                    .map(|(k, v)| (k.to_string(), v.to_string())),
            );

        if let Some(token) = source_config.get("token") {
            let handler = todo!("implement handler");
            let mut client =
                serenity::Client::builder(token, serenity::all::GatewayIntents::non_privileged())
                    .event_handler(handler.clone())
                    .await
                    .context("failed to create Discord client")?;
            // Spawn a background task for handling messages
            let task = tokio::spawn(async move {
                debug!("handling client start in task");
                if let Err(why) = client.start().await {
                    error!("Client error: {:?}", why);
                }
            });

            // TODO: cache the handler and client

            Ok(())
        } else {
            Err(anyhow::anyhow!(
                "token is required for discord authentication"
            ))
        }
    }

    async fn delete_link(&self, component_id: &str) -> Result<(), anyhow::Error> {
        debug!(component_id, "deleting link");
        let _ = self.handlers.write().await.remove(component_id);

        Ok(())
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let provider = DiscordProvider::new();
    let shutdown = run_provider(provider.clone(), "messaging-discord-provider")
        .await
        .context("failed to run provider")?;
    let connection = get_connection();
    serve(
        &connection.get_wrpc_client(connection.provider_key()),
        provider,
        shutdown,
    )
    .await
}

Though we don't have a message handler yet, this properly refactors the code to create Discord clients in response to links. We fetch the token from the configuration and use it to create an authenticated client. In order to implement the handler, we can create a separate file src/discord.rs to keep the code tidy:

rust

//! Discord specific logic using the [serenity] crate.

use std::{collections::HashMap, sync::Arc};

use serenity::{
    async_trait,
    model::{channel::Message, gateway::Ready},
    prelude::{Context as SerenityContext, *},
};
use wasmcloud_provider_sdk::get_connection;
use wit_bindgen_wrpc::tracing::{error, info, trace};

use crate::wasmcloud::messaging::types;

#[derive(Clone)]
pub struct DiscordHandler {
    /// The component ID that this handler is associated with
    pub component_id: String,
    /// A map of message IDs to Serenity context and message, so a component can
    /// reply asynchronously to a message
    pub messages: Arc<RwLock<HashMap<String, (SerenityContext, Message)>>>,
}

impl DiscordHandler {
    pub fn new(component_id: &str) -> Self {
        Self {
            component_id: component_id.to_string(),
            messages: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn message(&self, message_id: &str) -> Option<(SerenityContext, Message)> {
        self.messages.read().await.get(message_id).cloned()
    }
}

#[async_trait]
impl EventHandler for DiscordHandler {
    async fn message(&self, ctx: SerenityContext, msg: Message) {
        let discord_msg = types::BrokerMessage {
            subject: msg.channel_id.to_string(),
            reply_to: Some(msg.id.to_string()),
            body: msg.content.as_bytes().to_vec(),
        };

        let msg_id = msg.id.to_string();
        self.messages
            .write()
            .await
            .insert(msg.id.to_string(), (ctx, msg));

        // Spawns a task to remove the message from the map after 30 seconds, which should be
        // plenty of time for the component to reply.
        let messages = self.messages.clone();
        tokio::spawn(async move {
            let _ = tokio::time::sleep(std::time::Duration::from_secs(30)).await;
            messages.write().await.remove(&msg_id);
        });

        match crate::wasmcloud::messaging::handler::handle_message(
            &get_connection().get_wrpc_client(&self.component_id),
            &discord_msg,
        )
        .await
        {
            Ok(Ok(_)) => trace!("Component handled message successfully"),
            Ok(Err(e)) => error!(e, "Component failed to handle message"),
            Err(e) => error!(%e, "RPC error handling message"),
        }
    }

    async fn ready(&self, _: SerenityContext, ready: Ready) {
        info!(bot_name = ready.user.name, "Discord bot connected");
    }
}

Creating the handler required some reading into the EventHandler trait, but it wasn't difficult to implement. For each Message received from the Discord API, I simply restructured the message to fit the wasmcloud_messaging::BrokerMessage structure, and used the generated handle_message function to send that message to my component.

Finally, we can add the handler to our main file:

rust

#[derive(Default, Clone)]
struct DiscordProvider {
    /// A map of component ID to [DiscordHandler] which contains the Serenity client and message handlers
    handlers: Arc<RwLock<HashMap<String, DiscordHandler>>>,
    /// A map of component ID to the task handle for the client
    client_tasks: Arc<RwLock<HashMap<String, JoinHandle<()>>>>,
}

impl DiscordProvider {
    fn new() -> Self {
        Self {
            handlers: Arc::new(RwLock::new(HashMap::new())),
            client_tasks: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

impl Provider for DiscordProvider {
    async fn receive_link_config_as_source(
        &self,
        config: wasmcloud_provider_sdk::LinkConfig<'_>,
    ) -> Result<(), anyhow::Error> {
        debug!("receiving link configuration as source");
        let source_config: case_insensitive_hashmap::CaseInsensitiveHashMap<String> =
            case_insensitive_hashmap::CaseInsensitiveHashMap::from_iter(
                config
                    .config
                    .iter()
                    .map(|(k, v)| (k.to_string(), v.to_string())),
            );

        if let Some(token) = source_config.get("token") {
            let handler = DiscordHandler::new(config.target_id);
            let mut client =
                serenity::Client::builder(token, serenity::all::GatewayIntents::non_privileged())
                    .event_handler(handler.clone())
                    .await
                    .context("failed to create Discord client")?;
            let task = tokio::spawn(async move {
                debug!("handling client start in task");
                if let Err(why) = client.start().await {
                    error!("Client error: {:?}", why);
                }
            });

            self.handlers
                .write()
                .await
                .insert(config.target_id.to_string(), handler);
            self.client_tasks
                .write()
                .await
                .insert(config.target_id.to_string(), task);

            Ok(())
        } else {
            Err(anyhow::anyhow!(
                "token is required for discord authentication"
            ))
        }
    }
    
    // ...

Now our provider is fully set up to receive messages and forward them onto a component. What about when the component wants to send a message back? Given our example GoodJanet bot, we need to implement the publish function to respond to messages:

rust

async fn publish(
    &self,
    ctx: Option<Context>,
    msg: types::BrokerMessage,
) -> Result<Result<(), String>, anyhow::Error> {
    debug!("component publishing message as bot");
    let Some(Some(component_id)) = ctx.as_ref().map(|ctx| ctx.component.as_ref()) else {
        error!("missing component ID");
        return Ok(Err("missing component ID".to_string()));
    };

    let handlers = self.handlers.read().await;
    let component_handler = handlers.get(component_id).unwrap();

    let Some((ctx, original_msg)) = component_handler.message(&msg.subject).await else {
        error!(
            message_id = msg.subject,
            "component published message with unknown ID"
        );
        return Ok(Err("message not found for specified ID".to_string()));
    };

    let message_text = String::from_utf8_lossy(&msg.body);
    // This shouldn't really ever happen, but just in case this warning will help identify cases
    // where we need better string handling of incoming messages
    if message_text.contains(char::REPLACEMENT_CHARACTER) {
        warn!("message body is not valid UTF-8, may contain invalid characters");
    }
    original_msg
        .channel_id
        .say(&ctx.http, message_text)
        .await
        .map(|_| Ok(()))
        .map_err(|e| anyhow::anyhow!(e))
}

Now whenever a component calls wasmcloud:messaging/consumer.publish, we'll look up the original context of the message and send the message text back. We've added in a good bit of error handling here as well, just in case a component publishes with the wrong reply subject.

info

You can take a look at the completed provider in my wasmcloud-messaging-chatbot repository.

Running the provider

Putting it all together, we just need to create an application manifest for our bot and provider to deploy them to wasmCloud. We'll make sure to link the provider to the bot on the handler interface, supplying the token as configuration. Additionally, we'll link the bot to the provider in the other direction on the consumer interface to publish messages. This matches exactly what is shown in the component's WIT world.

wit

package wasmcloud:hello;

world hello {
  // Component calls function, link from component -> provider
  import wasmcloud:messaging/consumer@0.2.0;
  
  // Provider calls function, link from provider -> component
  export wasmcloud:messaging/handler@0.2.0;
}

yaml

apiVersion: core.oam.dev/v1beta1
kind: Application
metadata:
  name: good-janet
  annotations:
    version: v0.0.1
    description: "Good Janet Discord Bot"
spec:
  components:
    - name: janet
      type: component
      properties:
        image: file://./build/good_janet_s.wasm
      traits:
        # Govern the spread/scheduling of the component
        - type: spreadscaler
          properties:
            replicas: 1
        - type: link
          properties:
            target: discord
            namespace: wasmcloud
            package: messaging
            interfaces: [consumer]

    # Add a provider that implements `wasmcloud:messaging` using the Discord API
    - name: discord
      type: capability
      properties:
        image: file://../build/provider-messaging-discord.par.gz
      traits:
        - type: link
          properties:
            target: janet
            namespace: wasmcloud
            package: messaging
            interfaces: [handler]
            source_config:
              - name: janet-bot-token

In this manifest I reference configuration named janet-bot-token, but since this is a secret I'd prefer not to store it in plaintext alongside my application. Following this pattern, you simply need to put the configuration ahead of time with wash:

bash

wash config put janet-bot-token token=<bot-token>
wash app deploy ./wadm.yaml

Testing it out

Finally, after deploying our application, we can ping our Discord bot to see if it all worked:

Success!

Conclusion

This blog should serve as a guide for how you can start implementing a new capability provider, following a similar process to our create provider documentation for a different use case. If you have questions about creating custom providers or want to share what you've built, join the wasmCloud community Slack!