use std::fs;
use bevy_app::OnAppExitSystems;
use bevy_app::prelude::*;
use bevy_ecs::prelude::*;
use rig_core::effect::{Decider, EffectKind};
use rig_ecs::agent::Notice;
use rig_ecs::stage::{Decision, EffectRequest, Outcome, Settled, Started};
use rig_ecs::tasks::Wake;
use rig_harness_protocol::{RELOAD_EXIT_CODE, SessionId};
use super::session::SessionPaths;
pub const SIGNAL_EXIT_CODE: u8 = 130;
#[derive(Resource, Clone, Debug, PartialEq, Eq)]
pub enum ExitReason {
Quit,
Reload,
Switch(Option<SessionId>),
Signal,
Failed(u8),
}
impl ExitReason {
pub fn of(exit: &AppExit) -> Self {
match exit {
AppExit::Success => Self::Quit,
AppExit::Error(code) if code.get() == RELOAD_EXIT_CODE => Self::Reload,
AppExit::Error(code) if code.get() == SIGNAL_EXIT_CODE => Self::Signal,
AppExit::Error(code) => Self::Failed(code.get()),
}
}
pub fn exit(&self) -> AppExit {
match self {
Self::Quit => AppExit::Success,
Self::Reload | Self::Switch(_) => AppExit::from_code(RELOAD_EXIT_CODE),
Self::Signal => AppExit::from_code(SIGNAL_EXIT_CODE),
Self::Failed(code) => AppExit::from_code(*code),
}
}
}
#[derive(Event, Clone, Debug)]
pub struct Exit(pub ExitReason);
#[derive(Component, Clone, Debug)]
pub struct ExitRequest(pub ExitReason);
pub struct ExitPlugin;
impl Plugin for ExitPlugin {
fn build(&self, app: &mut App) {
app.add_observer(on_exit)
.add_observer(answer_started)
.add_observer(exit_settled)
.add_systems(
Last,
reason_from_exit
.before(OnAppExitSystems)
.run_if(on_message::<AppExit>),
);
}
}
fn on_exit(exit: On<Exit>, mut commands: Commands) {
let reason = exit.0.clone();
commands.queue(move |world: &mut World| {
if !world.contains_resource::<ExitReason>() {
if let ExitReason::Switch(target) = &reason
&& let Err(why) = name_switch(world, target.as_ref())
{
world.write_message(Notice::error(
None,
format!("Could not switch sessions: {why}."),
));
return;
}
world.insert_resource(reason.clone());
}
let kind = EffectKind::Custom {
kind: "exit".into(),
payload: format!("{reason:?}").into(),
};
world.spawn((
Name::new("exit"),
ExitRequest(reason),
EffectRequest {
kind,
handler: None,
},
));
if let Some(wake) = world.get_resource::<Wake>() {
wake.wake();
}
});
}
fn name_switch(world: &World, target: Option<&SessionId>) -> Result<(), String> {
match (
world.get_resource::<SessionPaths>(),
crate::harness_protocol::launcher(),
) {
(Some(paths), Some(_)) => {
let target = target.map(SessionId::as_str).unwrap_or_default();
fs::write(paths.switch(), target).map_err(|failure| failure.to_string())
}
_ => {
Err("switching sessions needs the rig launcher: start the agent with `rig`".to_owned())
}
}
}
fn answer_started(started: On<Add<Started>>, mut exits: Query<&mut Decision, With<ExitRequest>>) {
if let Ok(mut decision) = exits.get_mut(started.entity) {
decision.decide(Decider::plugin("rig-harness"), Outcome::custom("exited"));
}
}
fn exit_settled(
settled: On<Add<Settled>>,
requests: Query<&ExitRequest>,
reason: Option<Res<ExitReason>>,
mut exits: MessageWriter<AppExit>,
mut exited: Local<bool>,
) {
if let Ok(request) = requests.get(settled.entity)
&& !*exited
{
*exited = true;
exits.write(reason.as_deref().unwrap_or(&request.0).exit());
}
}
pub(super) fn reason_from_exit(
mut exits: MessageReader<AppExit>,
reason: Option<Res<ExitReason>>,
mut commands: Commands,
) {
let first = exits.read().next();
if let Some(exit) = first.filter(|_| reason.is_none()) {
commands.insert_resource(ExitReason::of(exit));
}
}
#[cfg(test)]
mod tests;