use nagisa_core::{
collect_into, CooldownStore, EnabledOverrides, EnabledSet, FlightStore, Handler, KillSwitch,
Matcher, Middleware, Rendezvous, Router, Rule, SleepState, Superusers, WaiterStore,
};
use nagisa_types::id::Uin;
use std::sync::Arc;
#[cfg(any(feature = "onebot", feature = "milky"))]
use {
nagisa_core::{run_dispatch, Bot, Service, ShutdownToken, Supervisor},
nagisa_types::error::Result,
nagisa_types::event::Event,
std::time::Duration,
tokio::sync::mpsc,
tokio::task::JoinSet,
};
#[cfg(any(feature = "onebot", feature = "milky"))]
const LOGIN_RETRIES: usize = 20;
#[cfg(any(feature = "onebot", feature = "milky"))]
const LOGIN_RETRY_DELAY: Duration = Duration::from_millis(150);
#[cfg(any(feature = "onebot", feature = "milky"))]
const EVENT_CHANNEL_CAP: usize = 256;
pub struct App {
router: Router,
enabled: Arc<EnabledSet>,
waiters: Arc<WaiterStore>,
rendezvous: Arc<Rendezvous<String, Uin>>,
cooldowns: Arc<CooldownStore>,
flight: Arc<FlightStore>,
kill_switch: Arc<KillSwitch>,
sleep: Arc<SleepState>,
#[cfg(any(feature = "onebot", feature = "milky"))]
services: Vec<Arc<dyn Service>>,
}
impl Default for App {
fn default() -> Self {
Self::new()
}
}
impl App {
pub fn new() -> Self {
let enabled = Arc::new(EnabledSet::new());
let waiters = Arc::new(WaiterStore::new());
let rendezvous = Arc::new(Rendezvous::<String, Uin>::default());
let cooldowns = Arc::new(CooldownStore::new());
let flight = Arc::new(FlightStore::new());
let kill_switch = Arc::new(KillSwitch::new());
let sleep = Arc::new(SleepState::new());
Self {
router: collect_into(Router::new())
.data_arc(Arc::clone(&enabled))
.data_arc(Arc::clone(&waiters))
.data_arc(Arc::clone(&rendezvous))
.data_arc(Arc::clone(&cooldowns))
.data_arc(Arc::clone(&flight))
.data_arc(Arc::clone(&kill_switch))
.data_arc(Arc::clone(&sleep)),
enabled,
waiters,
rendezvous,
cooldowns,
flight,
kill_switch,
sleep,
#[cfg(any(feature = "onebot", feature = "milky"))]
services: Vec::new(),
}
}
pub fn debug(mut self) -> Self {
self.router = self.router.data(nagisa_core::DevMode);
self
}
pub fn command<H, Args>(mut self, matcher: Matcher, handler: H) -> Self
where
H: Handler<Args>,
Args: 'static,
{
self.router = self.router.command(matcher, handler);
self
}
pub fn on<H, Args>(mut self, handler: H) -> Self
where
H: Handler<Args>,
Args: 'static,
{
self.router = self.router.on(handler);
self
}
pub fn on_top<H, Args>(mut self, handler: H) -> Self
where
H: Handler<Args>,
Args: 'static,
{
self.router = self.router.on_top(handler);
self
}
pub fn command_with<H, Args>(mut self, matcher: Matcher, gate: Rule, handler: H) -> Self
where
H: Handler<Args>,
Args: 'static,
{
self.router = self.router.command_with(matcher, gate, handler);
self
}
pub fn on_with<H, Args>(mut self, gate: Rule, handler: H) -> Self
where
H: Handler<Args>,
Args: 'static,
{
self.router = self.router.on_with(gate, handler);
self
}
pub fn layer<M: Middleware>(mut self, m: M) -> Self {
self.router = self.router.layer(m);
self
}
pub fn data<T: Send + Sync + 'static>(mut self, value: T) -> Self {
self.router = self.router.data(value);
self
}
pub fn superusers(self, ids: impl IntoIterator<Item = impl Into<crate::Uin>>) -> Self {
let set: std::collections::HashSet<crate::Uin> = ids.into_iter().map(Into::into).collect();
self.data(Superusers(set))
}
pub fn restore_switches(self, ov: EnabledOverrides) -> Self {
self.enabled.restore(ov);
self
}
pub fn enabled_handle(&self) -> Arc<EnabledSet> {
Arc::clone(&self.enabled)
}
pub fn waiter_store_handle(&self) -> Arc<WaiterStore> {
Arc::clone(&self.waiters)
}
pub fn rendezvous_handle(&self) -> Arc<Rendezvous<String, Uin>> {
Arc::clone(&self.rendezvous)
}
pub fn cooldown_store_handle(&self) -> Arc<CooldownStore> {
Arc::clone(&self.cooldowns)
}
pub fn flight_store_handle(&self) -> Arc<FlightStore> {
Arc::clone(&self.flight)
}
pub fn kill_switch_handle(&self) -> Arc<KillSwitch> {
Arc::clone(&self.kill_switch)
}
pub fn sleep_handle(&self) -> Arc<SleepState> {
Arc::clone(&self.sleep)
}
pub fn on_switch_change<F>(self, f: F) -> Self
where
F: Fn(&str, Option<nagisa_types::id::Peer>, bool) + Send + Sync + 'static,
{
self.enabled.on_change(f);
self
}
#[cfg(any(feature = "onebot", feature = "milky"))]
pub fn service(mut self, svc: Arc<dyn Service>) -> Self {
self.services.push(svc);
self
}
#[cfg(feature = "onebot")]
pub async fn run_onebot(
self,
cfg: nagisa_onebot::OneBotConfig,
shutdown: ShutdownToken,
) -> Result<()> {
let adapter = nagisa_onebot::OneBotAdapter::new(cfg);
self.run_with(adapter, shutdown).await
}
#[cfg(feature = "milky")]
pub async fn run_milky(
self,
cfg: nagisa_milky::MilkyConfig,
shutdown: ShutdownToken,
) -> Result<()> {
let adapter = Arc::new(nagisa_milky::MilkyAdapter::new(cfg)?);
self.run_with(adapter, shutdown).await
}
#[cfg(any(feature = "onebot", feature = "milky"))]
async fn run_with<A>(self, adapter: Arc<A>, shutdown: ShutdownToken) -> Result<()>
where
A: nagisa_core::EventSource + nagisa_core::adapter::Actions,
{
let router = Arc::new(self.router);
let (tx, rx) = mpsc::channel::<Event>(EVENT_CHANNEL_CAP);
let mut tasks: JoinSet<Result<()>> = JoinSet::new();
{
let source = Arc::clone(&adapter);
let sink = tx.clone();
let source_shutdown = shutdown.clone();
tasks.spawn(async move { source.run(sink, source_shutdown).await });
}
let (self_id, nickname) = resolve_self_id(
Arc::clone(&adapter) as Arc<dyn nagisa_core::adapter::ActionInvoker>,
&shutdown,
)
.await;
let bot = Bot::new(adapter, self_id);
{
let router = Arc::clone(&router);
let dispatch_shutdown = shutdown.clone();
tasks.spawn(async move {
run_dispatch(router, bot, rx, dispatch_shutdown).await;
Ok(())
});
}
if self_id.0 != 0 {
let _ = tx
.send(Event::Meta(nagisa_types::event::Meta::Ready { self_id, nickname }))
.await;
}
drop(tx);
if !self.services.is_empty() {
let mut supervisor = Supervisor::new();
for svc in self.services {
supervisor = supervisor.add(svc);
}
let svc_shutdown = shutdown.clone();
tasks.spawn(async move {
supervisor.run(svc_shutdown).await
});
}
run_until_done(tasks, shutdown).await
}
}
#[cfg(any(feature = "onebot", feature = "milky"))]
async fn resolve_self_id(
invoker: Arc<dyn nagisa_core::adapter::ActionInvoker>,
shutdown: &ShutdownToken,
) -> (Uin, String) {
for attempt in 0..LOGIN_RETRIES {
if shutdown.is_cancelled() {
return (Uin(0), String::new());
}
match invoker.get_login_info().await {
Ok((uin, nickname)) => {
tracing::info!(self_id = uin.0, nickname = %nickname, "resolved bot login info");
return (uin, nickname);
}
Err(e) => {
tracing::debug!(attempt, error = %e, "get_login_info not ready yet; retrying");
tokio::select! {
_ = shutdown.cancelled() => return (Uin(0), String::new()),
_ = tokio::time::sleep(LOGIN_RETRY_DELAY) => {}
}
}
}
}
tracing::warn!(
"could not resolve bot self_id via get_login_info; mention matching will be inert until reconnect"
);
(Uin(0), String::new())
}
#[cfg(any(feature = "onebot", feature = "milky"))]
async fn run_until_done(mut tasks: JoinSet<Result<()>>, shutdown: ShutdownToken) -> Result<()> {
let mut result: Result<()> = Ok(());
loop {
tokio::select! {
biased;
_ = shutdown.cancelled() => break,
joined = tasks.join_next() => {
match joined {
None => break, Some(Ok(Ok(()))) => continue, Some(Ok(Err(e))) => {
tracing::error!(error = %e, "app task returned error; shutting down");
result = Err(e);
break;
}
Some(Err(join_err)) => {
tracing::error!(error = %join_err, "app task panicked; shutting down");
break;
}
}
}
}
}
shutdown.cancel();
tasks.shutdown().await;
result
}