use std::sync::Arc;
pub use builder::*;
use crate::{
serenity_prelude::{self as serenity, TeamMemberRole},
BoxFuture,
};
mod builder;
pub struct Framework<U, E> {
user_data: std::sync::OnceLock<U>,
bot_id: std::sync::OnceLock<serenity::UserId>,
options: crate::FrameworkOptions<U, E>,
shard_manager: Option<Arc<serenity::ShardManager>>,
setup: std::sync::Mutex<
Option<
Box<
dyn Send
+ Sync
+ for<'a> FnOnce(
&'a serenity::Context,
&'a serenity::Ready,
&'a Self,
) -> BoxFuture<'a, Result<U, E>>,
>,
>,
>,
edit_tracker_purge_task: Option<tokio::task::JoinHandle<()>>,
}
impl<U, E> Framework<U, E> {
#[deprecated = "Please use Framework::builder instead"]
pub fn build() -> FrameworkBuilder<U, E> {
FrameworkBuilder::default()
}
pub fn builder() -> FrameworkBuilder<U, E> {
FrameworkBuilder::default()
}
pub fn new<F>(options: crate::FrameworkOptions<U, E>, setup: F) -> Self
where
F: Send
+ Sync
+ 'static
+ for<'a> FnOnce(
&'a serenity::Context,
&'a serenity::Ready,
&'a Self,
) -> BoxFuture<'a, Result<U, E>>,
U: Send + Sync + 'static,
E: Send + 'static,
{
Self {
user_data: std::sync::OnceLock::new(),
bot_id: std::sync::OnceLock::new(),
setup: std::sync::Mutex::new(Some(Box::new(setup))),
edit_tracker_purge_task: None,
shard_manager: None,
options,
}
}
pub fn options(&self) -> &crate::FrameworkOptions<U, E> {
&self.options
}
pub fn shard_manager(&self) -> &Arc<serenity::ShardManager> {
self.shard_manager
.as_ref()
.expect("framework should have started")
}
pub async fn user_data(&self) -> &U {
loop {
match self.user_data.get() {
Some(x) => break x,
None => tokio::time::sleep(std::time::Duration::from_millis(100)).await,
}
}
}
}
impl<U, E> Drop for Framework<U, E> {
fn drop(&mut self) {
if let Some(task) = &mut self.edit_tracker_purge_task {
task.abort()
}
}
}
#[serenity::async_trait]
impl<U: Send + Sync, E: Send + Sync> serenity::Framework for Framework<U, E> {
async fn init(&mut self, client: &serenity::Client) {
set_qualified_names(&mut self.options.commands);
message_content_intent_sanity_check(
&self.options.prefix_options,
client.shard_manager.intents(),
);
self.shard_manager = Some(client.shard_manager.clone());
if self.options.initialize_owners {
if let Err(e) = insert_owners_from_http(
&client.http,
&mut self.options.owners,
&self.options.initialized_team_roles,
)
.await
{
tracing::warn!("Failed to insert owners from HTTP: {e}");
}
}
if let Some(edit_tracker) = &self.options.prefix_options.edit_tracker {
self.edit_tracker_purge_task =
Some(spawn_edit_tracker_purge_task(edit_tracker.clone()));
}
}
async fn dispatch(&self, ctx: serenity::Context, event: serenity::FullEvent) {
raw_dispatch_event(self, ctx, event).await
}
}
async fn raw_dispatch_event<U, E>(
framework: &Framework<U, E>,
ctx: serenity::Context,
event: serenity::FullEvent,
) where
U: Send + Sync,
{
if let serenity::FullEvent::Ready { data_about_bot } = &event {
let _: Result<_, _> = framework.bot_id.set(data_about_bot.user.id);
let setup = Option::take(&mut *framework.setup.lock().unwrap());
if let Some(setup) = setup {
match setup(&ctx, data_about_bot, framework).await {
Ok(user_data) => {
let _: Result<_, _> = framework.user_data.set(user_data);
}
Err(error) => {
(framework.options.on_error)(crate::FrameworkError::Setup {
error,
framework,
data_about_bot,
ctx: &ctx,
})
.await
}
}
} else {
}
}
let user_data = framework.user_data().await;
let bot_id = *framework
.bot_id
.get()
.expect("bot ID not set even though we awaited Ready");
let framework = crate::FrameworkContext {
bot_id,
options: &framework.options,
user_data,
shard_manager: framework.shard_manager(),
};
crate::dispatch_event(framework, &ctx, event).await;
}
pub fn set_qualified_names<U, E>(commands: &mut [crate::Command<U, E>]) {
fn set_subcommand_qualified_names<U, E>(parents: &str, commands: &mut [crate::Command<U, E>]) {
for cmd in commands {
cmd.qualified_name = format!("{} {}", parents, cmd.name);
set_subcommand_qualified_names(&cmd.qualified_name, &mut cmd.subcommands);
}
}
for command in commands {
set_subcommand_qualified_names(&command.name, &mut command.subcommands);
}
}
fn message_content_intent_sanity_check<U, E>(
prefix_options: &crate::PrefixFrameworkOptions<U, E>,
intents: serenity::GatewayIntents,
) {
let is_prefix_configured = prefix_options.prefix.is_some()
|| prefix_options.dynamic_prefix.is_some()
|| prefix_options.stripped_dynamic_prefix.is_some();
let can_receive_message_content = intents.contains(serenity::GatewayIntents::MESSAGE_CONTENT);
if is_prefix_configured && !can_receive_message_content {
tracing::warn!(
"Warning: MESSAGE_CONTENT intent not set; prefix commands will not be received"
);
}
}
pub async fn insert_owners_from_http(
http: &serenity::Http,
owners: &mut std::collections::HashSet<serenity::UserId>,
initialized_teams: &Option<Vec<serenity::TeamMemberRole>>,
) -> Result<(), serenity::Error> {
let application_info = http.get_current_application_info().await?;
if let Some(owner) = application_info.owner {
owners.insert(owner.id);
}
if let Some(team) = application_info.team {
for member in team.members {
if member.membership_state != serenity::MembershipState::Accepted {
continue;
}
let Some(initialized_teams) = initialized_teams else {
if let TeamMemberRole::Admin | TeamMemberRole::Developer = member.role {
owners.insert(member.user.id);
}
continue;
};
if initialized_teams.contains(&member.role) {
owners.insert(member.user.id);
}
}
}
Ok(())
}
fn spawn_edit_tracker_purge_task(
edit_tracker: Arc<std::sync::RwLock<crate::EditTracker>>,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
loop {
edit_tracker.write().unwrap().purge();
tokio::time::sleep(std::time::Duration::from_secs(60)).await;
}
})
}