1use std::collections::BTreeSet;
9use std::fmt::{Debug, Write};
10use std::fs::{self, File, OpenOptions};
11use std::io::{self, Seek, SeekFrom};
12use std::ops::Deref;
13use std::path::{Path, PathBuf};
14use std::sync::{Arc, Mutex, PoisonError};
15
16use bevy_app::OnAppExitSystems;
17use bevy_app::prelude::*;
18use bevy_ecs::prelude::*;
19use bevy_log::tracing::field::{Field, Visit};
20use bevy_log::tracing::{Event, Level};
21use bevy_log::tracing_subscriber::layer::{Context, Layer};
22use bevy_log::tracing_subscriber::{Registry, fmt};
23use bevy_log::{BoxedFmtLayer, BoxedLayer, error, warn};
24use bevy_reflect::Reflect;
25use crossbeam_channel::{Receiver, Sender};
26use rig_harness_protocol::{Home, SessionDir, SessionId, StartReason};
27
28use rig_core::conversation::Author;
29use rig_ecs::WriteLogs;
30use rig_ecs::agent::{Notice, NoticeLevel};
31use rig_ecs::restore::SavedKeys;
32use rig_ecs::session_log::{SessionLog, SessionStore, open_session};
33use rig_memory::{
34 DirFiles, DirStore, Files, Rewrite, Rewritten, Session, SessionError, SessionMeta, Store,
35 rewrite, rewrite_file,
36};
37use rig_tools::Spill;
38
39use super::cli::{self, Cli};
40use super::exit::ExitReason;
41use super::launcher::{LauncherRun, refuse};
42use crate::describe::Describing;
43use crate::plugin_data::PluginData;
44use crate::services::{ProjectRoot, RigHome};
45use crate::timings::{self, SystemTimings, TimeSystems};
46use crate::workspace::ProjectWorkspace;
47
48#[derive(Resource, Clone, Debug)]
53pub struct SessionPaths(pub SessionDir);
54
55impl Deref for SessionPaths {
56 type Target = SessionDir;
57
58 fn deref(&self) -> &SessionDir {
59 &self.0
60 }
61}
62
63#[derive(Resource, Clone)]
80pub struct StoreOpener {
81 open: Arc<OpenStore>,
83 listed: Arc<BTreeSet<String>>,
85}
86
87type OpenStore = dyn Fn(&SessionDir) -> Result<Arc<dyn Store>, SessionError> + Send + Sync;
89
90impl StoreOpener {
91 pub fn new(
93 open: impl Fn(&SessionDir) -> Result<Arc<dyn Store>, SessionError> + Send + Sync + 'static,
94 ) -> Self {
95 Self {
96 open: Arc::new(open),
97 listed: Arc::default(),
98 }
99 }
100
101 pub fn open(&self, dir: &SessionDir) -> Result<Arc<dyn Store>, SessionError> {
108 (self.open)(dir)
109 }
110
111 pub fn relist(&self, store: &dyn Store, session: &Session) -> io::Result<()> {
122 let mut meta = SessionMeta::read(store).unwrap_or_default();
123 meta.relist(session, &self.listed);
124 meta.write(store)
125 }
126}
127
128impl Debug for StoreOpener {
129 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130 f.write_str("StoreOpener")
131 }
132}
133
134impl Default for StoreOpener {
135 fn default() -> Self {
136 Self::new(|dir| Ok(Arc::new(DirStore::open(dir.path())?)))
137 }
138}
139
140#[derive(Resource, Clone, Debug, PartialEq, Eq)]
148pub struct Scratch(pub Option<PathBuf>);
149
150impl Scratch {
151 pub fn spill(&self) -> Option<Spill> {
154 self.0.as_ref().map(|dir| Spill(dir.join("spill")))
155 }
156
157 pub fn images(&self) -> Option<PathBuf> {
160 self.0.as_ref().map(|dir| dir.join("images"))
161 }
162
163 pub fn log(&self) -> Option<PathBuf> {
165 self.0.as_ref().map(|dir| dir.join("agent.log"))
166 }
167}
168
169pub(crate) fn paths_in(home: &Home) -> SessionPaths {
172 let id = SessionId::from_env()
173 .ok()
174 .flatten()
175 .unwrap_or_else(SessionId::generate);
176 let dir = home.session(&id);
177 if crate::describe::describing().is_none() {
180 fs::create_dir_all(dir.path()).ok();
181 }
182 SessionPaths(dir)
183}
184
185fn session_paths(world: &mut World) -> SessionPaths {
188 let home = RigHome::of(world);
189 world
190 .get_resource_or_insert_with(|| paths_in(&home))
191 .clone()
192}
193
194#[derive(Resource, Clone, Copy, Debug, PartialEq, Eq)]
199pub struct AppStart(pub StartReason);
200
201#[derive(Resource)]
204struct Unmarked;
205
206#[derive(Resource)]
210struct Unsaved;
211
212pub struct SessionPlugin;
220
221impl Plugin for SessionPlugin {
222 fn build(&self, app: &mut App) {
223 std::panic::set_hook(Box::new(|info| {
224 error!("{info}\n{}", std::backtrace::Backtrace::capture());
225 }));
226 let paths = session_paths(app.world_mut());
227 if !app.world().contains_resource::<Scratch>() {
228 app.insert_resource(Scratch(Some(paths.path().to_owned())));
229 }
230 let first = !paths.started().exists();
231 if !app.world().contains_resource::<AppStart>() {
232 let reason = match StartReason::from_env() {
233 _ if first => StartReason::New,
234 None | Some(StartReason::New) => StartReason::Resumed,
235 Some(reason) => reason,
236 };
237 app.insert_resource(AppStart(reason));
238 }
239 if first {
240 app.insert_resource(Unmarked).add_systems(
241 Last,
242 mark_started
243 .run_if(resource_exists::<Unmarked>)
244 .before(discard_unsaved),
245 );
246 }
247 app.init_resource::<StoreOpener>()
248 .init_resource::<SystemTimings>()
249 .add_message::<SessionDeleted>()
250 .add_message::<SessionDeleteFailed>()
251 .add_message::<Notice>()
252 .add_observer(delete_session)
253 .add_systems(PreStartup, list_keys)
254 .add_systems(
255 Startup,
256 finish_deletes.run_if(not(resource_exists::<Describing>)),
257 )
258 .add_systems(First, timings::next_frame)
259 .add_systems(
260 Last,
261 mark_saved
262 .after(WriteLogs)
263 .before(discard_unsaved)
264 .run_if(resource_exists::<Unsaved>),
265 )
266 .add_systems(
267 Last,
268 discard_unsaved
269 .after(OnAppExitSystems)
270 .run_if(on_message::<AppExit>),
271 );
272 }
273
274 fn finish(&self, app: &mut App) {
281 if !app.world().contains_resource::<Describing>() {
282 match open_store(app.world_mut()) {
283 Ok(()) => {
284 let wanted = app.world().get_resource::<Scratch>().and_then(Scratch::log);
285 move_text_log(app.world_mut(), wanted);
286 }
287 Err(failure) => {
288 refuse(
289 app,
290 format!(
291 "the session cannot be opened, so nothing of it was changed: {failure}"
292 ),
293 );
294 }
295 }
296 }
297 insert_plugin_data(app.world_mut());
298 }
299}
300
301fn list_keys(keys: Option<Res<SavedKeys>>, stores: Option<ResMut<StoreOpener>>) {
304 if let Some(mut stores) = stores {
305 let listed = keys.map(|keys| keys.listed().clone());
306 stores.listed = Arc::new(listed.unwrap_or_default());
307 }
308}
309
310fn insert_plugin_data(world: &mut World) {
314 if world.contains_resource::<PluginData>() {
315 return;
316 }
317 let home = RigHome::of(world);
318 let session = world
319 .get_resource::<SessionStore>()
320 .map(|store| -> Arc<dyn Files> { store.0.clone() });
321 let root = world
322 .get_resource::<ProjectRoot>()
323 .cloned()
324 .unwrap_or_default();
325 let workspace = world.get_resource::<ProjectWorkspace>();
326 let host = workspace.and_then(|workspace| workspace.0.host());
327 let data = PluginData::new(&home, session, &root.0, host);
328 world.insert_resource(data);
329}
330
331fn open_store(world: &mut World) -> Result<(), SessionError> {
337 let store = match world.remove_resource::<SessionStore>() {
338 Some(SessionStore(store)) => store,
339 None => {
340 cli::choose(world);
341 let front = world.get_resource::<Cli>().and_then(Cli::chosen);
342 let keeps = !front.is_some_and(|front| front.sessionless);
343 let stores = front
344 .and_then(|front| front.session_store.as_ref())
345 .or_else(|| world.get_resource::<StoreOpener>());
346 let (Some(paths), Some(stores)) = (world.get_resource::<SessionPaths>(), stores) else {
347 return Ok(());
348 };
349 let store = stores.open(paths)?;
350 if keeps && !paths.is_saved() {
351 world.insert_resource(Unsaved);
352 }
353 store
354 }
355 };
356 open_session(world, store)
357}
358
359pub(crate) fn leave_session(world: &mut World) {
364 let elsewhere = world.get_resource::<LauncherRun>().map(|run| run.0.log());
365 move_text_log(world, elsewhere);
366 world.insert_resource(Scratch(None));
367 world.remove_resource::<Unmarked>();
368}
369
370fn mark_saved(log: Option<Res<SessionLog>>, paths: Res<SessionPaths>, mut commands: Commands) {
372 if !log.is_some_and(|log| log.has_written()) {
373 return;
374 }
375 commands.remove_resource::<Unsaved>();
376 if let Err(failure) = paths.mark_saved() {
377 warn!("could not mark the session saved: {failure}");
378 }
379}
380
381fn mark_started(paths: Res<SessionPaths>, mut commands: Commands) {
383 commands.remove_resource::<Unmarked>();
384 if let Err(failure) = fs::write(paths.started(), b"") {
385 error!("could not mark the session started: {failure}");
386 }
387}
388
389fn discard_unsaved(
400 mut exits: MessageReader<AppExit>,
401 (cli, start, reason): (
402 Option<Res<Cli>>,
403 Option<Res<AppStart>>,
404 Option<Res<ExitReason>>,
405 ),
406 paths: Res<SessionPaths>,
407 (run, scratch): (Option<Res<LauncherRun>>, Option<Res<Scratch>>),
408) {
409 let exits: Vec<&AppExit> = exits.read().collect();
410 let reload = match reason.as_deref() {
411 Some(reason) => *reason == ExitReason::Reload,
412 None => exits
413 .iter()
414 .any(|exit| ExitReason::of(exit) == ExitReason::Reload),
415 };
416 let discard = if cli.is_some_and(|cli| cli.sessionless()) {
417 start.is_none_or(|start| start.0 != StartReason::Resumed)
418 } else {
419 !reload && !paths.is_saved() && !paths.switch().exists()
420 };
421 if !discard {
422 return;
423 }
424 if !exits.iter().all(|exit| exit.is_success())
425 && let Some(run) = run
426 && let Some(log) = scratch.and_then(|scratch| scratch.log())
427 {
428 fs::rename(log, run.0.log()).ok();
429 }
430 fs::remove_dir_all(paths.path()).ok();
431}
432
433pub fn edit_stopped(
447 stores: &StoreOpener,
448 running: &SessionDir,
449 dir: &SessionDir,
450 edit: impl FnOnce(&mut Session) -> Result<(), String>,
451) -> Result<(), String> {
452 let _held = hold_stopped(running, dir, "edit its agents' conversations instead")?;
453 let store = stores.open(dir).map_err(|failure| failure.to_string())?;
454 let mut session = Session::read(&*store).map_err(|failure| failure.to_string())?;
455 edit(&mut session)?;
456 session
457 .save(&*store)
458 .map_err(|failure| failure.to_string())?;
459 relist(stores, &*store, &session);
460 Ok(())
461}
462
463fn relist(stores: &StoreOpener, store: &dyn Store, session: &Session) {
466 if let Err(failure) = stores.relist(store, session) {
467 warn!("could not write a stopped session's listing: {failure}");
468 }
469}
470
471fn hold_stopped(running: &SessionDir, dir: &SessionDir, this: &str) -> Result<File, String> {
475 if dir.path() == running.path() {
476 return Err(format!("it is this session: {this}"));
477 }
478 if !dir.path().is_dir() {
479 return Err("there is no such session".to_owned());
480 }
481 dir.hold()
482 .map_err(|failure| failure.to_string())?
483 .ok_or_else(|| "another rig runs it".to_owned())
484}
485
486pub fn rewrite_stopped(
521 stores: &StoreOpener,
522 running: &SessionDir,
523 dir: &SessionDir,
524 by: Author,
525 with: &dyn Rewrite,
526) -> Result<Rewritten, String> {
527 let _held = hold_stopped(running, dir, "only a stopped session is rewritten")?;
528 let store = stores.open(dir).map_err(|failure| failure.to_string())?;
529 let rewritten = rewrite(&*store, by, with).map_err(|failure| failure.to_string())?;
530 let scratch = DirFiles::new(dir.path());
531 map_spill(&scratch, with)
532 .and_then(|()| map_name(dir, with))
533 .and_then(|_| rewrite_file(&scratch, LOG, with))
534 .map_err(|failure| format!("rewriting its scratch files failed: {failure}"))?;
535 match Session::read(&*store) {
536 Ok(session) => relist(stores, &*store, &session),
537 Err(failure) => warn!("could not read a rewritten session to list it: {failure}"),
538 }
539 Ok(rewritten)
540}
541
542const LOG: &str = "agent.log";
544
545fn map_name(dir: &SessionDir, with: &dyn Rewrite) -> io::Result<bool> {
549 let Some(name) = dir.name() else {
550 return Ok(false);
551 };
552 let mut bytes = name.into_bytes();
553 if !with.file("name", &mut bytes) {
554 return Ok(false);
555 }
556 dir.set_name(Some(&String::from_utf8_lossy(&bytes)))?;
557 Ok(true)
558}
559
560fn map_spill(scratch: &DirFiles, with: &dyn Rewrite) -> io::Result<()> {
563 for name in scratch.files()? {
564 if name.starts_with("spill/") {
565 rewrite_file(scratch, &name, with)?;
566 }
567 }
568 Ok(())
569}
570
571#[derive(Event, Clone, Debug)]
579pub struct DeleteSession(pub SessionId);
580
581#[derive(Message, Clone, Debug)]
584pub struct SessionDeleted(pub SessionId);
585
586#[derive(Message, Clone, Debug)]
590pub struct SessionDeleteFailed {
591 pub id: SessionId,
593 pub why: String,
595}
596
597fn delete_session(
598 delete: On<DeleteSession>,
599 (paths, stores): (Res<SessionPaths>, Res<StoreOpener>),
600 home: Option<Res<RigHome>>,
601 mut deleted: MessageWriter<SessionDeleted>,
602 mut failed: MessageWriter<SessionDeleteFailed>,
603 mut notices: MessageWriter<Notice>,
604) {
605 let home = home.map_or_else(Home::from_env, |home| home.0.clone());
606 let session = home.session(&delete.0);
607 let refusal = if session.path() == paths.path() {
608 Some("it is this session".to_owned())
609 } else if !session.path().is_dir() {
610 Some("there is no such session".to_owned())
611 } else if session.runs_elsewhere() {
612 Some("another rig runs it".to_owned())
613 } else {
614 session
615 .mark_deleting()
616 .map_err(SessionError::from)
617 .and_then(|()| finish_delete(&session, &stores))
618 .err()
619 .map(|failure| failure.to_string())
620 };
621 match refusal {
622 None => {
623 deleted.write(SessionDeleted(delete.0.clone()));
624 }
625 Some(why) => {
626 notices.write(Notice::error(
627 None,
628 format!("Could not delete session {}: {why}.", delete.0),
629 ));
630 let id = delete.0.clone();
631 failed.write(SessionDeleteFailed { id, why });
632 }
633 }
634}
635
636fn finish_delete(session: &SessionDir, stores: &StoreOpener) -> Result<(), SessionError> {
640 match (stores.open)(session) {
641 Ok(store) => store.destroy()?,
642 Err(SessionError::OldLayout) => {}
643 Err(failure) => return Err(failure),
644 }
645 Ok(fs::remove_dir_all(session.path())?)
646}
647
648fn finish_deletes(
651 home: Option<Res<RigHome>>,
652 (paths, stores): (Res<SessionPaths>, Res<StoreOpener>),
653 mut deleted: MessageWriter<SessionDeleted>,
654 mut failed: MessageWriter<SessionDeleteFailed>,
655 mut notices: MessageWriter<Notice>,
656) {
657 let Some(home) = home else {
658 return;
659 };
660 let Ok(entries) = fs::read_dir(home.0.sessions()) else {
661 return;
662 };
663 for entry in entries.flatten() {
664 let session = SessionDir::at(entry.path());
665 if !session.is_deleting() || session.path() == paths.path() || session.runs_elsewhere() {
666 continue;
667 }
668 match (finish_delete(&session, &stores), session.id()) {
669 (Ok(()), Some(id)) => {
670 deleted.write(SessionDeleted(id));
671 }
672 (Ok(()), None) => {}
673 (Err(failure), id) => {
674 notices.write(Notice::warn(
675 None,
676 format!(
677 "Could not finish deleting the session in {}: {failure}.",
678 session.path().display()
679 ),
680 ));
681 if let Some(id) = id {
682 let why = failure.to_string();
683 failed.write(SessionDeleteFailed { id, why });
684 }
685 }
686 }
687 }
688}
689
690pub(crate) fn log_layer(app: &mut App) -> Option<BoxedFmtLayer> {
696 let at = session_paths(app.world_mut()).log();
697 let earlier = fs::metadata(&at).ok().map(|earlier| earlier.len());
698 let file = OpenOptions::new().create(true).append(true).open(&at).ok();
699 let log = TextLog {
700 at: file.is_some().then_some(at),
701 earlier,
702 file: Arc::new(Mutex::new(file)),
703 };
704 let writer = log.file.clone();
705 app.insert_resource(log);
706 let layer = fmt::Layer::default().with_ansi(false);
707 Some(Box::new(
708 layer.with_writer(move || LogWriter(writer.clone())),
709 ))
710}
711
712#[derive(Resource)]
714struct TextLog {
715 at: Option<PathBuf>,
716 earlier: Option<u64>,
719 file: Arc<Mutex<Option<File>>>,
720}
721
722struct LogWriter(Arc<Mutex<Option<File>>>);
724
725impl io::Write for LogWriter {
726 fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
727 let mut file = self.0.lock().unwrap_or_else(PoisonError::into_inner);
728 match file.as_mut() {
729 Some(file) => file.write(bytes),
730 None => Ok(bytes.len()),
731 }
732 }
733
734 fn flush(&mut self) -> io::Result<()> {
735 let mut file = self.0.lock().unwrap_or_else(PoisonError::into_inner);
736 file.as_mut().map_or(Ok(()), io::Write::flush)
737 }
738}
739
740fn move_text_log(world: &mut World, wanted: Option<PathBuf>) {
744 let Some(mut log) = world.get_resource_mut::<TextLog>() else {
745 return;
746 };
747 if log.at == wanted {
748 return;
749 }
750 let handle = log.file.clone();
751 let mut file = handle.lock().unwrap_or_else(PoisonError::into_inner);
752 let moved = match (&log.at, log.earlier, &wanted) {
753 (Some(at), None, Some(to)) => move_file(at, to),
754 (Some(at), None, None) => fs::remove_file(at),
755 (Some(at), Some(earlier), to) => take_tail(at, earlier, to.as_deref()),
756 (None, _, _) => Ok(()),
757 };
758 if let Err(failure) = moved {
759 drop(file);
761 error!("could not move the text log: {failure}");
762 return;
763 }
764 *file = wanted
765 .as_deref()
766 .and_then(|to| OpenOptions::new().create(true).append(true).open(to).ok());
767 drop(file);
768 log.at = wanted;
769 log.earlier = None;
770}
771
772fn take_tail(from: &Path, length: u64, to: Option<&Path>) -> io::Result<()> {
775 if let Some(to) = to {
776 let mut tail = File::open(from)?;
777 tail.seek(SeekFrom::Start(length))?;
778 if let Some(directory) = to.parent() {
779 fs::create_dir_all(directory)?;
780 }
781 let mut moved = OpenOptions::new().create(true).append(true).open(to)?;
782 io::copy(&mut tail, &mut moved)?;
783 }
784 OpenOptions::new().write(true).open(from)?.set_len(length)
785}
786
787fn move_file(from: &Path, to: &Path) -> io::Result<()> {
790 if let Some(directory) = to.parent() {
791 fs::create_dir_all(directory)?;
792 }
793 fs::rename(from, to).or_else(|_| {
794 fs::copy(from, to)?;
795 fs::remove_file(from)
796 })
797}
798
799#[derive(Reflect, Clone, Debug, Default)]
801pub struct Logged {
802 pub level: NoticeLevel,
804 pub target: String,
806 pub message: String,
808 pub agent: Option<String>,
810 pub plugin: Option<String>,
812}
813
814#[derive(Resource)]
817pub struct LoggedQueue(pub Receiver<Logged>);
818
819pub fn log_events(app: &mut App) -> Option<BoxedLayer> {
823 let (sender, receiver) = crossbeam_channel::bounded(256);
824 let timings = SystemTimings::default();
825 app.insert_resource(LoggedQueue(receiver))
826 .insert_resource(timings.clone());
827 Some(Box::new(PassOn(sender).and_then(TimeSystems(timings))))
828}
829
830struct PassOn(Sender<Logged>);
831
832impl Layer<Registry> for PassOn {
833 fn on_event(&self, event: &Event<'_>, _: Context<'_, Registry>) {
834 let level = *event.metadata().level();
835 if level <= Level::WARN {
836 let target = event.metadata().target().to_owned();
837 let level = if level == Level::ERROR {
838 NoticeLevel::Error
839 } else {
840 NoticeLevel::Warn
841 };
842 let mut logged = Logged {
843 level,
844 target,
845 ..Logged::default()
846 };
847 event.record(&mut logged);
848 self.0.try_send(logged).ok();
849 }
850 }
851}
852
853impl Visit for Logged {
856 fn record_str(&mut self, field: &Field, value: &str) {
857 match field.name() {
858 "log.target" => self.target = value.to_owned(),
859 _ => self.record_debug(field, &format_args!("{value}")),
860 }
861 }
862
863 fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
864 match field.name() {
865 "agent" => self.agent = Some(format!("{value:?}")),
866 "message" => self.message.insert_str(0, &format!("{value:?}")),
867 name if name.starts_with("log.") => {}
868 name => write!(self.message, " {name}={value:?}").unwrap_or_default(),
869 }
870 }
871}
872
873#[cfg(test)]
874mod tests;