use std::collections::BTreeSet;
use std::fmt::{Debug, Write};
use std::fs::{self, File, OpenOptions};
use std::io::{self, Seek, SeekFrom};
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, PoisonError};
use bevy_app::OnAppExitSystems;
use bevy_app::prelude::*;
use bevy_ecs::prelude::*;
use bevy_log::tracing::field::{Field, Visit};
use bevy_log::tracing::{Event, Level};
use bevy_log::tracing_subscriber::layer::{Context, Layer};
use bevy_log::tracing_subscriber::{Registry, fmt};
use bevy_log::{BoxedFmtLayer, BoxedLayer, error, warn};
use bevy_reflect::Reflect;
use crossbeam_channel::{Receiver, Sender};
use rig_harness_protocol::{Home, SessionDir, SessionId, StartReason};
use rig_core::conversation::Author;
use rig_ecs::WriteLogs;
use rig_ecs::agent::{Notice, NoticeLevel};
use rig_ecs::restore::SavedKeys;
use rig_ecs::session_log::{SessionLog, SessionStore, open_session};
use rig_memory::{
DirFiles, DirStore, Files, Rewrite, Rewritten, Session, SessionError, SessionMeta, Store,
rewrite, rewrite_file,
};
use rig_tools::Spill;
use super::cli::{self, Cli};
use super::exit::ExitReason;
use super::launcher::{LauncherRun, refuse};
use crate::describe::Describing;
use crate::plugin_data::PluginData;
use crate::services::{ProjectRoot, RigHome};
use crate::timings::{self, SystemTimings, TimeSystems};
use crate::workspace::ProjectWorkspace;
#[derive(Resource, Clone, Debug)]
pub struct SessionPaths(pub SessionDir);
impl Deref for SessionPaths {
type Target = SessionDir;
fn deref(&self) -> &SessionDir {
&self.0
}
}
#[derive(Resource, Clone)]
pub struct StoreOpener {
open: Arc<OpenStore>,
listed: Arc<BTreeSet<String>>,
}
type OpenStore = dyn Fn(&SessionDir) -> Result<Arc<dyn Store>, SessionError> + Send + Sync;
impl StoreOpener {
pub fn new(
open: impl Fn(&SessionDir) -> Result<Arc<dyn Store>, SessionError> + Send + Sync + 'static,
) -> Self {
Self {
open: Arc::new(open),
listed: Arc::default(),
}
}
pub fn open(&self, dir: &SessionDir) -> Result<Arc<dyn Store>, SessionError> {
(self.open)(dir)
}
pub fn relist(&self, store: &dyn Store, session: &Session) -> io::Result<()> {
let mut meta = SessionMeta::read(store).unwrap_or_default();
meta.relist(session, &self.listed);
meta.write(store)
}
}
impl Debug for StoreOpener {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("StoreOpener")
}
}
impl Default for StoreOpener {
fn default() -> Self {
Self::new(|dir| Ok(Arc::new(DirStore::open(dir.path())?)))
}
}
#[derive(Resource, Clone, Debug, PartialEq, Eq)]
pub struct Scratch(pub Option<PathBuf>);
impl Scratch {
pub fn spill(&self) -> Option<Spill> {
self.0.as_ref().map(|dir| Spill(dir.join("spill")))
}
pub fn images(&self) -> Option<PathBuf> {
self.0.as_ref().map(|dir| dir.join("images"))
}
pub fn log(&self) -> Option<PathBuf> {
self.0.as_ref().map(|dir| dir.join("agent.log"))
}
}
pub(crate) fn paths_in(home: &Home) -> SessionPaths {
let id = SessionId::from_env()
.ok()
.flatten()
.unwrap_or_else(SessionId::generate);
let dir = home.session(&id);
if crate::describe::describing().is_none() {
fs::create_dir_all(dir.path()).ok();
}
SessionPaths(dir)
}
fn session_paths(world: &mut World) -> SessionPaths {
let home = RigHome::of(world);
world
.get_resource_or_insert_with(|| paths_in(&home))
.clone()
}
#[derive(Resource, Clone, Copy, Debug, PartialEq, Eq)]
pub struct AppStart(pub StartReason);
#[derive(Resource)]
struct Unmarked;
#[derive(Resource)]
struct Unsaved;
pub struct SessionPlugin;
impl Plugin for SessionPlugin {
fn build(&self, app: &mut App) {
std::panic::set_hook(Box::new(|info| {
error!("{info}\n{}", std::backtrace::Backtrace::capture());
}));
let paths = session_paths(app.world_mut());
if !app.world().contains_resource::<Scratch>() {
app.insert_resource(Scratch(Some(paths.path().to_owned())));
}
let first = !paths.started().exists();
if !app.world().contains_resource::<AppStart>() {
let reason = match StartReason::from_env() {
_ if first => StartReason::New,
None | Some(StartReason::New) => StartReason::Resumed,
Some(reason) => reason,
};
app.insert_resource(AppStart(reason));
}
if first {
app.insert_resource(Unmarked).add_systems(
Last,
mark_started
.run_if(resource_exists::<Unmarked>)
.before(discard_unsaved),
);
}
app.init_resource::<StoreOpener>()
.init_resource::<SystemTimings>()
.add_message::<SessionDeleted>()
.add_message::<SessionDeleteFailed>()
.add_message::<Notice>()
.add_observer(delete_session)
.add_systems(PreStartup, list_keys)
.add_systems(
Startup,
finish_deletes.run_if(not(resource_exists::<Describing>)),
)
.add_systems(First, timings::next_frame)
.add_systems(
Last,
mark_saved
.after(WriteLogs)
.before(discard_unsaved)
.run_if(resource_exists::<Unsaved>),
)
.add_systems(
Last,
discard_unsaved
.after(OnAppExitSystems)
.run_if(on_message::<AppExit>),
);
}
fn finish(&self, app: &mut App) {
if !app.world().contains_resource::<Describing>() {
match open_store(app.world_mut()) {
Ok(()) => {
let wanted = app.world().get_resource::<Scratch>().and_then(Scratch::log);
move_text_log(app.world_mut(), wanted);
}
Err(failure) => {
refuse(
app,
format!(
"the session cannot be opened, so nothing of it was changed: {failure}"
),
);
}
}
}
insert_plugin_data(app.world_mut());
}
}
fn list_keys(keys: Option<Res<SavedKeys>>, stores: Option<ResMut<StoreOpener>>) {
if let Some(mut stores) = stores {
let listed = keys.map(|keys| keys.listed().clone());
stores.listed = Arc::new(listed.unwrap_or_default());
}
}
fn insert_plugin_data(world: &mut World) {
if world.contains_resource::<PluginData>() {
return;
}
let home = RigHome::of(world);
let session = world
.get_resource::<SessionStore>()
.map(|store| -> Arc<dyn Files> { store.0.clone() });
let root = world
.get_resource::<ProjectRoot>()
.cloned()
.unwrap_or_default();
let workspace = world.get_resource::<ProjectWorkspace>();
let host = workspace.and_then(|workspace| workspace.0.host());
let data = PluginData::new(&home, session, &root.0, host);
world.insert_resource(data);
}
fn open_store(world: &mut World) -> Result<(), SessionError> {
let store = match world.remove_resource::<SessionStore>() {
Some(SessionStore(store)) => store,
None => {
cli::choose(world);
let front = world.get_resource::<Cli>().and_then(Cli::chosen);
let keeps = !front.is_some_and(|front| front.sessionless);
let stores = front
.and_then(|front| front.session_store.as_ref())
.or_else(|| world.get_resource::<StoreOpener>());
let (Some(paths), Some(stores)) = (world.get_resource::<SessionPaths>(), stores) else {
return Ok(());
};
let store = stores.open(paths)?;
if keeps && !paths.is_saved() {
world.insert_resource(Unsaved);
}
store
}
};
open_session(world, store)
}
pub(crate) fn leave_session(world: &mut World) {
let elsewhere = world.get_resource::<LauncherRun>().map(|run| run.0.log());
move_text_log(world, elsewhere);
world.insert_resource(Scratch(None));
world.remove_resource::<Unmarked>();
}
fn mark_saved(log: Option<Res<SessionLog>>, paths: Res<SessionPaths>, mut commands: Commands) {
if !log.is_some_and(|log| log.has_written()) {
return;
}
commands.remove_resource::<Unsaved>();
if let Err(failure) = paths.mark_saved() {
warn!("could not mark the session saved: {failure}");
}
}
fn mark_started(paths: Res<SessionPaths>, mut commands: Commands) {
commands.remove_resource::<Unmarked>();
if let Err(failure) = fs::write(paths.started(), b"") {
error!("could not mark the session started: {failure}");
}
}
fn discard_unsaved(
mut exits: MessageReader<AppExit>,
(cli, start, reason): (
Option<Res<Cli>>,
Option<Res<AppStart>>,
Option<Res<ExitReason>>,
),
paths: Res<SessionPaths>,
(run, scratch): (Option<Res<LauncherRun>>, Option<Res<Scratch>>),
) {
let exits: Vec<&AppExit> = exits.read().collect();
let reload = match reason.as_deref() {
Some(reason) => *reason == ExitReason::Reload,
None => exits
.iter()
.any(|exit| ExitReason::of(exit) == ExitReason::Reload),
};
let discard = if cli.is_some_and(|cli| cli.sessionless()) {
start.is_none_or(|start| start.0 != StartReason::Resumed)
} else {
!reload && !paths.is_saved() && !paths.switch().exists()
};
if !discard {
return;
}
if !exits.iter().all(|exit| exit.is_success())
&& let Some(run) = run
&& let Some(log) = scratch.and_then(|scratch| scratch.log())
{
fs::rename(log, run.0.log()).ok();
}
fs::remove_dir_all(paths.path()).ok();
}
pub fn edit_stopped(
stores: &StoreOpener,
running: &SessionDir,
dir: &SessionDir,
edit: impl FnOnce(&mut Session) -> Result<(), String>,
) -> Result<(), String> {
let _held = hold_stopped(running, dir, "edit its agents' conversations instead")?;
let store = stores.open(dir).map_err(|failure| failure.to_string())?;
let mut session = Session::read(&*store).map_err(|failure| failure.to_string())?;
edit(&mut session)?;
session
.save(&*store)
.map_err(|failure| failure.to_string())?;
relist(stores, &*store, &session);
Ok(())
}
fn relist(stores: &StoreOpener, store: &dyn Store, session: &Session) {
if let Err(failure) = stores.relist(store, session) {
warn!("could not write a stopped session's listing: {failure}");
}
}
fn hold_stopped(running: &SessionDir, dir: &SessionDir, this: &str) -> Result<File, String> {
if dir.path() == running.path() {
return Err(format!("it is this session: {this}"));
}
if !dir.path().is_dir() {
return Err("there is no such session".to_owned());
}
dir.hold()
.map_err(|failure| failure.to_string())?
.ok_or_else(|| "another rig runs it".to_owned())
}
pub fn rewrite_stopped(
stores: &StoreOpener,
running: &SessionDir,
dir: &SessionDir,
by: Author,
with: &dyn Rewrite,
) -> Result<Rewritten, String> {
let _held = hold_stopped(running, dir, "only a stopped session is rewritten")?;
let store = stores.open(dir).map_err(|failure| failure.to_string())?;
let rewritten = rewrite(&*store, by, with).map_err(|failure| failure.to_string())?;
let scratch = DirFiles::new(dir.path());
map_spill(&scratch, with)
.and_then(|()| map_name(dir, with))
.and_then(|_| rewrite_file(&scratch, LOG, with))
.map_err(|failure| format!("rewriting its scratch files failed: {failure}"))?;
match Session::read(&*store) {
Ok(session) => relist(stores, &*store, &session),
Err(failure) => warn!("could not read a rewritten session to list it: {failure}"),
}
Ok(rewritten)
}
const LOG: &str = "agent.log";
fn map_name(dir: &SessionDir, with: &dyn Rewrite) -> io::Result<bool> {
let Some(name) = dir.name() else {
return Ok(false);
};
let mut bytes = name.into_bytes();
if !with.file("name", &mut bytes) {
return Ok(false);
}
dir.set_name(Some(&String::from_utf8_lossy(&bytes)))?;
Ok(true)
}
fn map_spill(scratch: &DirFiles, with: &dyn Rewrite) -> io::Result<()> {
for name in scratch.files()? {
if name.starts_with("spill/") {
rewrite_file(scratch, &name, with)?;
}
}
Ok(())
}
#[derive(Event, Clone, Debug)]
pub struct DeleteSession(pub SessionId);
#[derive(Message, Clone, Debug)]
pub struct SessionDeleted(pub SessionId);
#[derive(Message, Clone, Debug)]
pub struct SessionDeleteFailed {
pub id: SessionId,
pub why: String,
}
fn delete_session(
delete: On<DeleteSession>,
(paths, stores): (Res<SessionPaths>, Res<StoreOpener>),
home: Option<Res<RigHome>>,
mut deleted: MessageWriter<SessionDeleted>,
mut failed: MessageWriter<SessionDeleteFailed>,
mut notices: MessageWriter<Notice>,
) {
let home = home.map_or_else(Home::from_env, |home| home.0.clone());
let session = home.session(&delete.0);
let refusal = if session.path() == paths.path() {
Some("it is this session".to_owned())
} else if !session.path().is_dir() {
Some("there is no such session".to_owned())
} else if session.runs_elsewhere() {
Some("another rig runs it".to_owned())
} else {
session
.mark_deleting()
.map_err(SessionError::from)
.and_then(|()| finish_delete(&session, &stores))
.err()
.map(|failure| failure.to_string())
};
match refusal {
None => {
deleted.write(SessionDeleted(delete.0.clone()));
}
Some(why) => {
notices.write(Notice::error(
None,
format!("Could not delete session {}: {why}.", delete.0),
));
let id = delete.0.clone();
failed.write(SessionDeleteFailed { id, why });
}
}
}
fn finish_delete(session: &SessionDir, stores: &StoreOpener) -> Result<(), SessionError> {
match (stores.open)(session) {
Ok(store) => store.destroy()?,
Err(SessionError::OldLayout) => {}
Err(failure) => return Err(failure),
}
Ok(fs::remove_dir_all(session.path())?)
}
fn finish_deletes(
home: Option<Res<RigHome>>,
(paths, stores): (Res<SessionPaths>, Res<StoreOpener>),
mut deleted: MessageWriter<SessionDeleted>,
mut failed: MessageWriter<SessionDeleteFailed>,
mut notices: MessageWriter<Notice>,
) {
let Some(home) = home else {
return;
};
let Ok(entries) = fs::read_dir(home.0.sessions()) else {
return;
};
for entry in entries.flatten() {
let session = SessionDir::at(entry.path());
if !session.is_deleting() || session.path() == paths.path() || session.runs_elsewhere() {
continue;
}
match (finish_delete(&session, &stores), session.id()) {
(Ok(()), Some(id)) => {
deleted.write(SessionDeleted(id));
}
(Ok(()), None) => {}
(Err(failure), id) => {
notices.write(Notice::warn(
None,
format!(
"Could not finish deleting the session in {}: {failure}.",
session.path().display()
),
));
if let Some(id) = id {
let why = failure.to_string();
failed.write(SessionDeleteFailed { id, why });
}
}
}
}
}
pub(crate) fn log_layer(app: &mut App) -> Option<BoxedFmtLayer> {
let at = session_paths(app.world_mut()).log();
let earlier = fs::metadata(&at).ok().map(|earlier| earlier.len());
let file = OpenOptions::new().create(true).append(true).open(&at).ok();
let log = TextLog {
at: file.is_some().then_some(at),
earlier,
file: Arc::new(Mutex::new(file)),
};
let writer = log.file.clone();
app.insert_resource(log);
let layer = fmt::Layer::default().with_ansi(false);
Some(Box::new(
layer.with_writer(move || LogWriter(writer.clone())),
))
}
#[derive(Resource)]
struct TextLog {
at: Option<PathBuf>,
earlier: Option<u64>,
file: Arc<Mutex<Option<File>>>,
}
struct LogWriter(Arc<Mutex<Option<File>>>);
impl io::Write for LogWriter {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
let mut file = self.0.lock().unwrap_or_else(PoisonError::into_inner);
match file.as_mut() {
Some(file) => file.write(bytes),
None => Ok(bytes.len()),
}
}
fn flush(&mut self) -> io::Result<()> {
let mut file = self.0.lock().unwrap_or_else(PoisonError::into_inner);
file.as_mut().map_or(Ok(()), io::Write::flush)
}
}
fn move_text_log(world: &mut World, wanted: Option<PathBuf>) {
let Some(mut log) = world.get_resource_mut::<TextLog>() else {
return;
};
if log.at == wanted {
return;
}
let handle = log.file.clone();
let mut file = handle.lock().unwrap_or_else(PoisonError::into_inner);
let moved = match (&log.at, log.earlier, &wanted) {
(Some(at), None, Some(to)) => move_file(at, to),
(Some(at), None, None) => fs::remove_file(at),
(Some(at), Some(earlier), to) => take_tail(at, earlier, to.as_deref()),
(None, _, _) => Ok(()),
};
if let Err(failure) = moved {
drop(file);
error!("could not move the text log: {failure}");
return;
}
*file = wanted
.as_deref()
.and_then(|to| OpenOptions::new().create(true).append(true).open(to).ok());
drop(file);
log.at = wanted;
log.earlier = None;
}
fn take_tail(from: &Path, length: u64, to: Option<&Path>) -> io::Result<()> {
if let Some(to) = to {
let mut tail = File::open(from)?;
tail.seek(SeekFrom::Start(length))?;
if let Some(directory) = to.parent() {
fs::create_dir_all(directory)?;
}
let mut moved = OpenOptions::new().create(true).append(true).open(to)?;
io::copy(&mut tail, &mut moved)?;
}
OpenOptions::new().write(true).open(from)?.set_len(length)
}
fn move_file(from: &Path, to: &Path) -> io::Result<()> {
if let Some(directory) = to.parent() {
fs::create_dir_all(directory)?;
}
fs::rename(from, to).or_else(|_| {
fs::copy(from, to)?;
fs::remove_file(from)
})
}
#[derive(Reflect, Clone, Debug, Default)]
pub struct Logged {
pub level: NoticeLevel,
pub target: String,
pub message: String,
pub agent: Option<String>,
pub plugin: Option<String>,
}
#[derive(Resource)]
pub struct LoggedQueue(pub Receiver<Logged>);
pub fn log_events(app: &mut App) -> Option<BoxedLayer> {
let (sender, receiver) = crossbeam_channel::bounded(256);
let timings = SystemTimings::default();
app.insert_resource(LoggedQueue(receiver))
.insert_resource(timings.clone());
Some(Box::new(PassOn(sender).and_then(TimeSystems(timings))))
}
struct PassOn(Sender<Logged>);
impl Layer<Registry> for PassOn {
fn on_event(&self, event: &Event<'_>, _: Context<'_, Registry>) {
let level = *event.metadata().level();
if level <= Level::WARN {
let target = event.metadata().target().to_owned();
let level = if level == Level::ERROR {
NoticeLevel::Error
} else {
NoticeLevel::Warn
};
let mut logged = Logged {
level,
target,
..Logged::default()
};
event.record(&mut logged);
self.0.try_send(logged).ok();
}
}
}
impl Visit for Logged {
fn record_str(&mut self, field: &Field, value: &str) {
match field.name() {
"log.target" => self.target = value.to_owned(),
_ => self.record_debug(field, &format_args!("{value}")),
}
}
fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
match field.name() {
"agent" => self.agent = Some(format!("{value:?}")),
"message" => self.message.insert_str(0, &format!("{value:?}")),
name if name.starts_with("log.") => {}
name => write!(self.message, " {name}={value:?}").unwrap_or_default(),
}
}
}
#[cfg(test)]
mod tests;