Skip to main content

rig_ecs/
lib.rs

1//! rig-ecs, the kernel of the rig coding agent: agents, their turns and
2//! their calls are Bevy entities, and the calls run on Bevy's task pools.
3//! Every call goes through one effect dispatch path, recorded as it settles
4//! ([`effects`]); plugins register the tools. The kernel saves every change
5//! of an agent's [`Conversation`](agent::Conversation) to the
6//! [`session_log::SessionStore`] the app inserts. Time is Bevy's: a retried
7//! model call waits on a delayed command, and plugins that animate or poll
8//! run on `on_real_timer`, not on threads of their own.
9//!
10//! The kernel depends on no view and no file system: the app fills in what
11//! it needs, such as the store (any [`rig_memory::Store`]) and the
12//! [`model::Models`].
13//!
14//! ```no_run
15//! use rig_ecs::session_log::SessionStore;
16//! use rig_ecs::prelude::*;
17//! use rig_memory::MemoryStore;
18//!
19//! let mut app = App::new();
20//! app.insert_resource(SessionStore::new(MemoryStore::default()))
21//!     .add_plugins((AgentPlugin, SessionLogPlugin));
22//! ```
23
24pub mod agent;
25pub mod effects;
26pub mod fork;
27pub mod inbox;
28pub mod model;
29pub mod prompt;
30pub mod restore;
31pub mod retry;
32pub mod session_log;
33pub mod stage;
34pub mod tasks;
35pub mod tools;
36pub mod turn;
37
38/// What a plugin needs: Bevy's app, ECS, reflection and time preludes (with
39/// the time run conditions and `Disabled`), the kernel's components,
40/// requests and stages, the tool registry, rig-core's conversation types and
41/// rig-memory's store. `Message` is Bevy's message trait, a conversation's
42/// message is `message::Message`, and `Conversation` is the agent's
43/// component.
44pub mod prelude {
45    pub use bevy_app::prelude::*;
46    pub use bevy_ecs::entity_disabling::Disabled;
47    pub use bevy_ecs::prelude::*;
48    pub use bevy_reflect::prelude::*;
49    pub use bevy_time::common_conditions::{on_real_timer, on_timer};
50    pub use bevy_time::prelude::*;
51
52    pub use crate::AgentPlugin;
53    pub use crate::agent::{
54        ActiveTurn, Agent, AgentId, CacheKey, CallOf, Calls, Conversation, EffectParent, LastUsage,
55        Notice, NoticeLevel, Partial, Retry, Spawned, SpawnedBy, StopTurn, SystemPrompt,
56        ToolAccess, Turn, TurnOf, TurnOutcome,
57    };
58    pub use crate::fork::{AgentForked, ForkAgent};
59    pub use crate::inbox::{
60        Attachment, Deliver, DeliveredIn, DeliveryMode, Incoming, MessagePart, MessageRecord,
61        Recalled, Unsent,
62    };
63    pub use crate::model::{
64        Connection, ModelChoice, Models, ReasoningSetting, SetModel, SetReasoning,
65    };
66    pub use crate::prompt::{AgentSections, PromptSection, SectionOf};
67    pub use crate::restore::{Interrupted, Restored, SavedKeys};
68    pub use crate::retry::Verdict;
69    pub use crate::session_log::{
70        ConversationChanged, ReflectSaved, Retire, SessionLogPlugin, SessionRestored,
71    };
72    pub use crate::stage::{
73        Admitted, AgentLoop, AgentPass, AgentWork, Decision, Effect, EffectRecorded, EffectRequest,
74        ForAgent, GateSystems, HoldOf, Holds, IfStopped, IncomingMessages, ModelCalls, Outcome,
75        PluginEffects, PrepareSystems, Published, ReadyToSettle, ReadyToStart, ReviewSystems,
76        Settled, Stage, Started, SupersededBy, Supersedes, ToolCalls, Waiting,
77        despawn_open_children, follow_successor,
78    };
79    pub use crate::tasks::{Done, PollTasks, Running, Wake};
80    pub use crate::tools::{
81        AppToolsExt, CallHandler, DeclaredFootprint, Footprint, ToolCalled, ToolDef, ToolOptions,
82        ToolRequest, WorldToolsExt,
83    };
84    pub use crate::turn::{
85        Adapter, Adapters, Backoff, DroppedOptions, ModelRequest, ModelStarts, Recovery,
86        RetryVerdict, SentTo, SideCall, StreamDelta, StreamTiming, TurnCall,
87    };
88    pub use rig_core::conversation::{
89        Aside, AsideId, Author, AuthorKind, Change, ChangeRecord, Entry, EntryRef, Halt, MessageId,
90        RequestId, STOPPED, Spliced, TurnEnd, TurnRecord,
91    };
92    pub use rig_core::effect::Decider;
93    pub use rig_core::message::{self, AssistantContent, AssistantMessage, ToolCall, UserContent};
94    pub use rig_core::now_ms;
95    pub use rig_core::streaming::{Part, PartKind, StreamEvent};
96    pub use rig_core::tool::{
97        Scope, Tool, ToolContext, ToolExecutionError, ToolResult, Touches, args_schema,
98    };
99    pub use rig_memory::{
100        Files, LogWriteError, Repair, Rewrite, Rewritten, SavedAgent, Session, Stamp, Store,
101        Strings, as_of, history,
102    };
103}
104
105use bevy_app::MainScheduleOrder;
106use bevy_app::prelude::*;
107use bevy_ecs::prelude::*;
108use bevy_ecs::schedule::ScheduleLabel;
109use bevy_ecs::schedule::SingleThreadedExecutor;
110use bevy_time::TimePlugin;
111
112use agent::{Agent, Notice};
113use session_log::{SaveConversations, SavedOne, SessionLog, save_conversations, save_now};
114use tasks::{PollTasks, Wake, poll_tasks, run_settle_frames};
115
116/// The system in `Last`, on exit, that stops the running turns and leaves
117/// them for the restart, unended; before [`WriteLogs`], both in Bevy's
118/// `OnAppExitSystems`.
119#[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
120pub struct StopTurnsOnExit;
121
122/// The systems in `Last` that queue what the frame changed for the session
123/// log: every conversation changed and every
124/// [`ReflectSaved`](session_log::ReflectSaved) value, after
125/// [`StopTurnsOnExit`] and before [`WriteLogs`]. A plugin's `Last` system that
126/// changes saved state goes `.before(SaveSystems)`: one after it is saved a
127/// frame late, and on the exit frame never.
128#[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
129pub struct SaveSystems;
130
131/// The systems in `Last` that write the frame's log records to the
132/// [`session_log::SessionStore`], and those a plugin adds to write its own
133/// logs, such as the effect log: such a system says `.in_set(WriteLogs)` and
134/// nothing else. It runs after [`SaveSystems`] and the turns stopped on
135/// exit, so it writes what they left too. A failed write stops the logging
136/// and sends a [`Notice`] here, at the end of the frame, which a system that
137/// runs before it reads the next frame and, on the exit frame, never: on
138/// exit, read [`SessionLog::failure`] or go `.after(WriteLogs)`.
139#[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
140pub struct WriteLogs;
141
142/// The kernel: agents, their turn loop, effects and the tool registry, and
143/// the saving of every conversation change (announced as
144/// [`session_log::ConversationChanged`]).
145///
146/// It spawns one agent at startup when the restored session has none, and
147/// logs to the app's [`session_log::SessionStore`] once
148/// [`session_log::SessionLogPlugin`] restored the session. It adds Bevy's
149/// `TimePlugin` unless the app has one, for the clock a retried call waits
150/// on, and then runs `First`, `PreUpdate` and `Last` single-threaded. It sets
151/// no error handler.
152pub struct AgentPlugin;
153
154impl Plugin for AgentPlugin {
155    fn build(&self, app: &mut App) {
156        if !app.is_plugin_added::<TimePlugin>() {
157            // An app without a clock of its own is a small headless one, not
158            // a game: its `First` and `PreUpdate` hold little but the clock's
159            // systems, and its `Last` little but the session log's, too little
160            // for Bevy's multi-threaded executor, which would wake the task
161            // pool for them every frame (a frame that saves a conversation
162            // paid more in thread handoffs than in saving).
163            let single = |schedule: &mut Schedule| {
164                schedule.set_executor(SingleThreadedExecutor::new());
165            };
166            app.add_plugins(TimePlugin)
167                .edit_schedule(First, single)
168                .edit_schedule(PreUpdate, single)
169                .edit_schedule(Last, single);
170        }
171        // A conversation changed before `Update` is saved before it runs.
172        if let Some(mut order) = app.world_mut().get_resource_mut::<MainScheduleOrder>()
173            && let Some(at) = order
174                .labels
175                .iter()
176                .position(|label| *label == Update.intern())
177        {
178            order.labels.insert(at, SaveConversations.intern());
179        }
180        app.init_resource::<effects::Effects>()
181            .init_resource::<SessionLog>()
182            .init_resource::<SavedOne>()
183            .edit_schedule(SaveConversations, |schedule| {
184                schedule.set_executor(SingleThreadedExecutor::new());
185            })
186            .add_systems(SaveConversations, save_conversations)
187            .init_resource::<Wake>()
188            .init_resource::<model::Models>()
189            .add_message::<Notice>()
190            .add_message::<session_log::ConversationChanged>()
191            .add_message::<inbox::Recalled>()
192            // Notices and commits are written by systems that never meet
193            // on one effect, and their order within a pass means nothing.
194            .allow_ambiguous_resource::<Messages<Notice>>()
195            .allow_ambiguous_resource::<Messages<session_log::ConversationChanged>>()
196            .configure_sets(
197                Last,
198                (StopTurnsOnExit, WriteLogs)
199                    .chain()
200                    .in_set(bevy_app::OnAppExitSystems),
201            )
202            .configure_sets(Last, SaveSystems.after(StopTurnsOnExit).before(WriteLogs))
203            .add_systems(Startup, spawn_first_agent)
204            // A reply that came in this frame has its calls started in it.
205            .configure_sets(Update, PollTasks.before(stage::AgentLoop))
206            .add_systems(Update, poll_tasks.in_set(PollTasks))
207            .add_systems(
208                PreUpdate,
209                model::reconnect.run_if(resource_changed::<model::Models>),
210            )
211            .add_systems(
212                Last,
213                (
214                    run_settle_frames,
215                    turn::stop_turns_on_exit
216                        .in_set(StopTurnsOnExit)
217                        .run_if(on_message::<AppExit>),
218                    save_now.map(drop).in_set(SaveSystems),
219                ),
220            )
221            .add_observer(session_log::retire)
222            .add_observer(fork::fork_agent)
223            .add_observer(inbox::on_deliver)
224            .add_observer(turn::on_retry)
225            .add_observer(turn::on_stop_turn)
226            .add_observer(inbox::unsent_on_turn_end)
227            .add_observer(inbox::unmark)
228            .add_observer(model::on_set_model)
229            .add_observer(model::connect)
230            .add_observer(model::check_effort)
231            .add_observer(model::on_set_reasoning);
232        stage::build(app);
233    }
234}
235
236fn spawn_first_agent(agents: Query<(), With<Agent>>, mut commands: Commands) {
237    if agents.is_empty() {
238        commands.spawn((Name::new("agent"), Agent));
239    }
240}