use std::borrow::Cow;
use std::cell::{Cell, RefCell};
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use bevy_app::prelude::*;
use bevy_ecs::error::{BevyError, ErrorContext, ErrorHandler, FallbackErrorHandler};
use bevy_ecs::lifecycle::HookContext;
use bevy_ecs::prelude::*;
use bevy_ecs::query::QueryData;
use bevy_ecs::schedule::{
InternedSystemSet, LogLevel, ScheduleBuildSettings, ScheduleLabel, SingleThreadedExecutor,
};
use bevy_ecs::world::DeferredWorld;
use bevy_log::tracing::Instrument;
use bevy_log::{info_span, warn};
use bevy_reflect::prelude::*;
use bevy_tasks::{AsyncComputeTaskPool, TaskPool};
use bevy_time::{Real, Time, Virtual};
use rig_core::completion::CompletionResponse;
use rig_core::effect::{
self, Decider, Disposition, EffectId, EffectKind, EffectRecord, HandlerKey, tool_key,
};
use rig_core::error::{ErrorKind, ErrorReport};
use rig_core::serve::{DispatchContext, Recorder, Reply, cancelled};
use rig_core::tool::{ToolExecutionError, ToolResult, ToolResultContext};
use rig_core::{millis, now_ms};
use web_time::SystemTime;
use super::agent::{AgentId, CallOf, Conversation, Notice, Turn, TurnOutcome};
use super::effects::{Effects, Handler, Ran};
use super::inbox::{Incoming, accept, carry_on_accepted, message_kind, settle_unsent};
use super::restore::{Interrupted, answer_interrupted, mark_started};
use super::session_log::{Drain, SessionLog, save_now};
use super::tasks::{Running, Wake, arrive};
use super::tools::{
CallHandler, Footprint, ToolRequest, ToolStart, ToolStarts, complete_calls, compute_footprints,
redact_results, refuse_calls,
};
use super::turn::{
Adapters, DroppedOptions, ModelRequest, ModelStarts, SentTo, StreamDelta, TurnCall, back_off,
build_requests, carry_on_called, commit_replies, commit_turn_calls, end_turns, feed,
judge_failures, keeps_settled, model_pool, stop_calls, stream_partials, stream_reply,
wait_in_order,
};
#[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
pub struct AgentPass;
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct AgentLoop;
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Stage {
Arrive,
Review,
Commit,
Admit,
Prepare,
Gate,
Start,
}
impl Stage {
const ALL: [Self; 7] = [
Self::Arrive,
Self::Review,
Self::Commit,
Self::Admit,
Self::Prepare,
Self::Gate,
Self::Start,
];
}
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum PrepareSystems {
Input,
Compose,
Edit,
Fit,
}
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ReviewSystems {
Edit,
Act,
}
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum GateSystems {
FromRecord,
After,
Fallback,
}
#[derive(SystemSet, Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct Enter(Stage);
#[derive(Component, Debug)]
#[component(on_add = take_id, on_remove = record_unsettled)]
#[require(Decision)]
pub struct Effect {
id: EffectId,
pub parent: Option<EffectId>,
pub purpose: Option<Cow<'static, str>>,
answered: Option<Disposition>,
raw: Option<Result<effect::Answered, ErrorReport>>,
decided_by: Option<Decider>,
}
impl Default for Effect {
fn default() -> Self {
Self::with_parent_id(None)
}
}
impl Effect {
pub fn with_parent_id(parent: Option<EffectId>) -> Self {
Self {
id: EffectId::from_raw(0),
parent,
purpose: None,
answered: None,
raw: None,
decided_by: None,
}
}
#[must_use]
pub fn with_purpose(mut self, purpose: impl Into<Cow<'static, str>>) -> Self {
self.purpose = Some(purpose.into());
self
}
pub fn id(&self) -> EffectId {
self.id
}
pub(crate) fn ran(&self) -> bool {
self.answered == Some(Disposition::Ran)
}
}
#[derive(Component, Clone)]
#[require(Effect)]
pub struct EffectRequest {
pub kind: EffectKind,
pub handler: Option<Handler>,
}
#[derive(Component, Reflect, Debug)]
#[reflect(Component)]
#[relationship(relationship_target = AgentWork)]
pub struct ForAgent(pub Entity);
#[derive(Component, Reflect, Debug)]
#[reflect(Component)]
#[relationship_target(relationship = ForAgent, linked_spawn)]
pub struct AgentWork(Vec<Entity>);
#[derive(Component, Reflect, Debug, Default)]
#[component(storage = "SparseSet", on_add = moved::<Admitted>)]
#[reflect(Component, Default)]
pub struct Admitted;
#[derive(Component, Clone, Copy, Debug)]
#[component(on_add = moved::<Started>)]
pub struct Started {
pub at: SystemTime,
}
#[derive(Component, Clone, Debug)]
#[component(on_insert = answered)]
pub struct Outcome(pub Result<effect::Answered, ErrorReport>);
impl Outcome {
pub fn refused(why: impl Into<String>) -> Self {
Self(Err(ErrorReport::new(ErrorKind::Denied, why).refused()))
}
pub fn completion(&self) -> Option<&CompletionResponse> {
match &self.0 {
Ok(effect::Answered::Completion(response)) => Some(response),
_ => None,
}
}
pub fn tool_result(&self) -> ToolResult {
match &self.0 {
Ok(effect::Answered::ToolResult { result }) => result.clone(),
Ok(other) => ToolResult::failed(ToolExecutionError::other(format!(
"the tool answered with a {} outcome",
other.family()
))),
Err(report) => ToolResult::failed(report.clone().into()),
}
}
pub fn custom(payload: impl Into<serde_json::Value>) -> Self {
Self(Ok(effect::Answered::Custom {
payload: payload.into(),
}))
}
}
impl From<ToolResult> for Outcome {
fn from(result: ToolResult) -> Self {
Self(Ok(effect::Answered::ToolResult { result }))
}
}
impl From<Result<effect::Answered, ErrorReport>> for Outcome {
fn from(outcome: Result<effect::Answered, ErrorReport>) -> Self {
Self(outcome)
}
}
#[derive(Component, Debug, Default)]
pub struct Decision(Option<(Outcome, Decider)>);
impl Decision {
pub fn decide(&mut self, by: Decider, outcome: impl Into<Outcome>) -> bool {
if self.0.is_some() {
return false;
}
self.0 = Some((outcome.into(), by));
true
}
pub fn get(&self) -> Option<&Outcome> {
self.0.as_ref().map(|(outcome, _)| outcome)
}
pub fn by(&self) -> Option<&Decider> {
self.0.as_ref().map(|(_, by)| by)
}
pub fn is_decided(&self) -> bool {
self.0.is_some()
}
}
#[derive(Component, Clone, Debug)]
pub struct IfStopped(pub Outcome);
#[derive(Component, Clone, Debug)]
pub struct Published(pub ToolResultContext);
#[derive(Component, Reflect, Debug)]
#[component(on_remove = moved::<HoldOf>)]
#[reflect(Component)]
#[relationship(relationship_target = Holds)]
pub struct HoldOf(pub Entity);
#[derive(Component, Reflect, Debug)]
#[reflect(Component)]
#[relationship_target(relationship = HoldOf, linked_spawn)]
pub struct Holds(Vec<Entity>);
#[derive(Component, Reflect, Debug)]
#[component(on_add = moved::<SupersededBy>)]
#[reflect(Component)]
#[relationship(relationship_target = Supersedes)]
pub struct SupersededBy(pub Entity);
impl SupersededBy {
pub fn successor(&self) -> Entity {
self.0
}
}
#[derive(Component, Reflect, Debug)]
#[reflect(Component)]
#[relationship_target(relationship = SupersededBy)]
pub struct Supersedes(Vec<Entity>);
#[derive(Component, Reflect, Debug, Default)]
#[component(storage = "SparseSet", on_add = moved::<Settled>)]
#[reflect(Component, Default)]
pub struct Settled;
pub fn follow_successor<M: Component>(
superseded: On<Add<SupersededBy>>,
owned: Query<(&SupersededBy, Option<&HoldOf>), With<M>>,
mut commands: Commands,
) {
let old = superseded.entity;
let Ok((by, hold)) = owned.get(old) else {
return;
};
let new = by.successor();
let held = hold.map(|&HoldOf(held)| held);
commands.queue(move |world: &mut World| {
let Ok(mut to) = world.get_entity_mut(new) else {
return;
};
if let Some(held) = held {
to.insert(HoldOf(held));
}
let Ok(mut from) = world.get_entity_mut(old) else {
return;
};
let marker = from.take::<M>();
if held.is_some() {
from.remove::<HoldOf>();
}
if let (Some(marker), Ok(mut to)) = (marker, world.get_entity_mut(new)) {
to.insert(marker);
}
});
}
pub fn despawn_open_children<M: Component>(
settled: On<Add<Settled>>,
children: Query<&Children>,
open: Query<(), (With<M>, Without<Settled>)>,
mut commands: Commands,
) {
let children = children.get(settled.entity).into_iter().flatten();
for &child in children.filter(|&&child| open.contains(child)) {
commands.entity(child).despawn();
}
}
#[derive(Message, Clone, Debug)]
pub struct EffectRecorded(pub EffectRecord);
#[derive(Component, Reflect, Debug, Default)]
#[component(storage = "SparseSet", on_remove = moved::<Waiting>)]
#[reflect(Component, Default)]
pub struct Waiting;
pub type ReadyToStart = (
With<Admitted>,
Without<Waiting>,
Without<Started>,
Without<Outcome>,
Without<Holds>,
Without<Settled>,
Without<SupersededBy>,
Without<Interrupted>,
);
pub type ReadyToSettle = (With<Outcome>, Without<Holds>, Without<Settled>);
pub type ToolCalls = (
With<ToolRequest>,
Without<ModelRequest>,
Without<Incoming>,
Without<EffectRequest>,
);
pub type ModelCalls = (
With<ModelRequest>,
Without<ToolRequest>,
Without<Incoming>,
Without<EffectRequest>,
);
pub type IncomingMessages = (
With<Incoming>,
Without<ToolRequest>,
Without<ModelRequest>,
Without<EffectRequest>,
);
pub type PluginEffects = (
With<EffectRequest>,
Without<ToolRequest>,
Without<ModelRequest>,
Without<Incoming>,
);
const MAX_PASSES: usize = 16;
#[derive(Resource, Default)]
struct Passes {
moved: u64,
quiet_at: u64,
last: Option<(&'static str, Entity)>,
noticed: HashSet<String>,
}
thread_local! {
static STAGE: Cell<Option<Stage>> = const { Cell::new(None) };
static FAILED: RefCell<Vec<Failure>> = const { RefCell::new(Vec::new()) };
static APP_HANDLER: Cell<Option<ErrorHandler>> = const { Cell::new(None) };
}
pub(crate) fn build(app: &mut App) {
app.init_resource::<Passes>()
.add_message::<EffectRecorded>()
.add_message::<StreamDelta>()
.edit_schedule(AgentPass, |pass| {
pass.set_executor(SingleThreadedExecutor::new());
if cfg!(debug_assertions) {
pass.set_build_settings(ScheduleBuildSettings {
ambiguity_detection: LogLevel::Warn,
..Default::default()
});
}
})
.configure_sets(
AgentPass,
(
Enter(Stage::Arrive),
Stage::Arrive,
Enter(Stage::Review),
Stage::Review,
Enter(Stage::Commit),
Stage::Commit,
Enter(Stage::Admit),
Stage::Admit,
Enter(Stage::Prepare),
Stage::Prepare,
Enter(Stage::Gate),
GateSystems::FromRecord,
Stage::Gate,
GateSystems::After,
GateSystems::Fallback,
Enter(Stage::Start),
Stage::Start,
)
.chain(),
)
.configure_sets(
AgentPass,
(
(
PrepareSystems::Input,
PrepareSystems::Compose,
PrepareSystems::Edit,
PrepareSystems::Fit,
)
.chain()
.in_set(Stage::Prepare),
(ReviewSystems::Edit, ReviewSystems::Act)
.chain()
.in_set(Stage::Review),
),
)
.add_systems(
AgentPass,
redact_results.in_set(Enter(Stage::Commit)),
)
.configure_sets(
AgentPass,
(
Stage::Review,
Stage::Commit,
Stage::Admit,
Stage::Prepare,
GateSystems::FromRecord,
Stage::Gate,
GateSystems::After,
GateSystems::Fallback,
Stage::Start,
)
.distributive_run_if(moved_since_quiet),
)
.add_systems(
AgentPass,
(
(arrive, stream_partials, judge_failures)
.chain()
.in_set(Stage::Arrive),
(
stop_calls,
back_off.run_if(passed(Stage::Review)),
settle,
commit_replies,
end_turns.run_if(passed(Stage::Review)),
commit_turn_calls,
)
.chain()
.in_set(Stage::Commit),
(complete_calls, admit, carry_on_called, build_requests)
.chain()
.in_set(Stage::Admit),
(
answer_unprepared,
compute_footprints.run_if(moved_since_quiet),
wait_in_order.run_if(moved_since_quiet),
)
.chain()
.in_set(Enter(Stage::Gate)),
refuse_calls
.after(GateSystems::FromRecord)
.before(Stage::Gate)
.run_if(moved_since_quiet),
(answer_interrupted, settle_unsent)
.chain()
.in_set(GateSystems::Fallback),
apply_decisions.in_set(Enter(Stage::Start)),
(accept.run_if(passed(Stage::Gate)), carry_on_accepted, start)
.chain()
.in_set(Stage::Start),
),
)
.add_systems(
Update,
run_passes
.in_set(AgentLoop)
.run_if(any_match_filter::<Or<(With<Effect>, With<Turn>)>>),
);
for stage in Stage::ALL {
app.add_systems(AgentPass, entering(stage).in_set(Enter(stage)));
}
let mut after: Option<InternedSystemSet> = None;
for stage in Stage::ALL {
let save = save_in_pass.before(Enter(stage));
match after {
Some(after) => app.add_systems(AgentPass, save.after(after)),
None => app.add_systems(AgentPass, save),
};
after = Some(match stage {
Stage::Gate => GateSystems::Fallback.intern(),
stage => stage.intern(),
});
}
app.add_systems(AgentPass, save_in_pass.after(Stage::Start));
}
fn save_in_pass(world: &mut World) {
if let Some(agent) = save_now(world)
&& let Some(mut passes) = world.get_resource_mut::<Passes>()
{
saved_some(&mut passes, agent);
}
}
fn saved_some(passes: &mut Passes, agent: Entity) {
passes.moved = passes.moved.wrapping_add(1);
passes.last = Some((std::any::type_name::<Conversation>(), agent));
}
fn entering(stage: Stage) -> impl Fn() + Clone + Send + Sync + 'static {
move || STAGE.set(Some(stage))
}
fn moved_since_quiet(passes: Res<Passes>) -> bool {
passes.moved != passes.quiet_at
}
fn run_passes(world: &mut World, settled: &mut QueryState<Entity, With<Settled>>) {
let done: Vec<Entity> = settled
.iter(world)
.filter(|&entity| !keeps_settled(world, entity))
.collect();
for entity in done {
world.try_despawn(entity).ok();
}
let mut clocks = None;
let app_handler = world.get_resource::<FallbackErrorHandler>().copied();
APP_HANDLER.set(Some(world.fallback_error_handler()));
world.insert_resource(FallbackErrorHandler(fail_closed));
let quiet = world.try_schedule_scope(AgentPass, |world, pass| {
(0..MAX_PASSES).any(|n| {
if n == 1 {
clocks = Some((
zero_step::<()>(world),
zero_step::<Real>(world),
zero_step::<Virtual>(world),
));
}
FAILED.with_borrow_mut(Vec::clear);
let before = moves(world);
pass.run(world);
moves(world) == before
})
});
STAGE.set(None);
FAILED.with_borrow_mut(Vec::clear);
APP_HANDLER.set(None);
match app_handler {
Some(handler) => world.insert_resource(handler),
None => {
world.remove_resource::<FallbackErrorHandler>();
}
}
if let Some((time, real, virtual_time)) = clocks {
put_back(world, time);
put_back(world, real);
put_back(world, virtual_time);
}
let Some(mut passes) = world.get_resource_mut::<Passes>() else {
return;
};
match quiet {
Ok(true) => passes.quiet_at = passes.moved,
Ok(false) => match passes.last {
Some((what, entity)) => warn!(
"the agent loop still moved after {MAX_PASSES} passes; the last change: {what} \
on {entity}"
),
None => warn!("the agent loop still moved after {MAX_PASSES} passes"),
},
Err(_) => {}
}
}
fn moves(world: &World) -> u64 {
world
.get_resource::<Passes>()
.map_or(0, |passes| passes.moved)
}
fn zero_step<T: Default + Copy + Send + Sync + 'static>(world: &mut World) -> Option<Time<T>> {
let mut time = world.get_resource_mut::<Time<T>>()?;
let was = *time;
time.advance_by(Duration::ZERO);
Some(was)
}
fn put_back<T: Default + Copy + Send + Sync + 'static>(world: &mut World, was: Option<Time<T>>) {
if let (Some(was), Some(mut time)) = (was, world.get_resource_mut::<Time<T>>()) {
*time = was;
}
}
fn fail_closed(error: BevyError, context: ErrorContext) {
let in_system = matches!(
context,
ErrorContext::System { .. } | ErrorContext::RunCondition { .. }
);
if in_system && let Some(stage) = STAGE.get() {
let message = error.to_string();
let failure = Failure {
stage,
system: context.to_string(),
error: message.lines().next().unwrap_or_default().to_owned(),
};
FAILED.with_borrow_mut(|failed| failed.push(failure));
}
let handler = APP_HANDLER
.get()
.unwrap_or(FallbackErrorHandler::default().0);
handler(error, context);
}
#[derive(Clone)]
pub(crate) struct Failure {
stage: Stage,
system: String,
error: String,
}
pub(crate) fn failed_in(stage: Stage) -> Option<Failure> {
FAILED.with_borrow(|failed| failed.iter().find(|failed| failed.stage == stage).cloned())
}
fn passed(stage: Stage) -> impl Fn() -> bool + Clone {
move || failed_in(stage).is_none()
}
fn closed(passes: &mut Passes, failure: Failure, what: &str) -> Option<Notice> {
let Failure {
stage,
system,
error,
} = failure;
let text = format!("{system} in {stage:?} ({error}): {what}.");
passes
.noticed
.insert(format!("{stage:?} {system}"))
.then(|| Notice::error(None, text))
}
pub(crate) fn moved<C: Component>(mut world: DeferredWorld, context: HookContext) {
if let Some(mut passes) = world.get_resource_mut::<Passes>() {
passes.moved = passes.moved.wrapping_add(1);
passes.last = Some((std::any::type_name::<C>(), context.entity));
}
}
fn take_id(mut world: DeferredWorld, context: HookContext) {
let entity = context.entity;
let id = world.get_resource::<Effects>().map(Effects::next_id);
let parent = world
.get::<ChildOf>(entity)
.and_then(|of| world.get::<Effect>(of.parent()))
.map(Effect::id);
if let Some(mut effect) = world.get_mut::<Effect>(entity) {
if let Some(id) = id {
effect.id = id;
}
if effect.parent.is_none() {
effect.parent = parent;
}
}
moved::<Effect>(world, context);
}
fn record_unsettled(mut world: DeferredWorld, context: HookContext) {
if world.get::<Settled>(context.entity).is_some() {
return;
}
if let Some(record) = assemble(&world, context.entity, Some(Disposition::Cancelled)) {
world.write_message(EffectRecorded(record));
}
}
fn answered(mut world: DeferredWorld, context: HookContext) {
let entity = context.entity;
let started = world.get::<Started>(entity).is_some();
let running = world.get::<Running>(entity).is_some();
let first = world
.get::<Effect>(entity)
.filter(|effect| effect.answered.is_none())
.and_then(|_| world.get::<Outcome>(entity))
.map(|outcome| outcome.0.clone());
if let (Some(first), Some(mut effect)) = (first, world.get_mut::<Effect>(entity)) {
effect.answered = Some(match (started, running) {
(false, _) => Disposition::AnsweredBeforeStart,
(true, true) => Disposition::AnsweredWhileRunning,
(true, false) => Disposition::Ran,
});
effect.raw = Some(first);
}
moved::<Outcome>(world, context);
}
fn apply_decisions(
mut decided: Query<
(Entity, &mut Effect, &Decision),
(Changed<Decision>, Without<Outcome>, Without<Settled>),
>,
mut commands: Commands,
) {
let mut decided: Vec<_> = decided
.iter_mut()
.filter_map(|(entity, effect, decision)| Some((entity, effect, decision.0.as_ref()?)))
.collect();
decided.sort_unstable_by_key(|(_, effect, _)| effect.id());
for (entity, mut effect, (outcome, by)) in decided {
effect.decided_by = Some(by.clone());
commands.entity(entity).try_insert(outcome.clone());
}
}
fn admit(
new: Query<Entity, (With<Effect>, Without<Admitted>, Without<Settled>)>,
mut commands: Commands,
) {
for entity in &new {
commands.entity(entity).try_insert(Admitted);
}
}
fn answer_unprepared(
admitted: Query<Entity, Added<Admitted>>,
mut passes: ResMut<Passes>,
mut commands: Commands,
mut notices: MessageWriter<Notice>,
) {
let Some(failure) = failed_in(Stage::Prepare) else {
return;
};
let why = format!("{} in Prepare: {}", failure.system, failure.error);
for entity in &admitted {
let failed = ErrorReport::new(ErrorKind::Internal, why.clone());
commands
.entity(entity)
.try_insert_if_new(Outcome(Err(failed)));
}
let what = "what it admitted is answered with an error";
notices.write_batch(closed(&mut passes, failure, what));
}
type Saving<'w, 's> = ParamSet<
'w,
's,
(
Drain<'static, 'static>,
Query<'static, 'static, &'static mut Conversation>,
),
>;
fn start(
mut ready: Query<StartItem, ReadyToStart>,
agents: Query<&AgentId>,
turns: Query<&TurnOutcome>,
(tools, models): (ToolStarts, ModelStarts),
effects: Res<Effects>,
wake: Res<Wake>,
(log, mut saving): (Res<SessionLog>, Saving),
mut passes: ResMut<Passes>,
mut commands: Commands,
mut notices: MessageWriter<Notice>,
) {
if let Some(failure) = failed_in(Stage::Gate) {
if !ready.is_empty() {
let what = "nothing started this pass";
notices.write_batch(closed(&mut passes, failure, what));
}
return;
}
let at = SystemTime::now();
let stopping = |item: &StartItemReadOnlyItem| {
let ending = item.of.and_then(|of| turns.get(of.0).ok());
ending.and_then(TurnOutcome::stops).is_some()
};
let changing: Vec<_> = ready
.iter()
.filter(|item| !item.message && !item.decision.is_decided() && !stopping(item))
.filter(|item| !item.footprint.is_some_and(|it| it.0.reads_only()))
.filter_map(|item| Some((item.agent?.0, item.tool?.call.id.clone())))
.collect();
let mut unsaved = None;
if !changing.is_empty() {
let mut conversations = saving.p1();
for (agent, call) in &changing {
if let Ok(mut conversation) = conversations.get_mut(*agent) {
mark_started(&mut conversation, call);
}
}
if let Some(agent) = saving.p0().save() {
saved_some(&mut passes, agent);
}
unsaved = log.flush().err().map(|failure| {
let why = format!(
"Not run: the session can no longer be saved ({failure}). Nothing changed."
);
ErrorReport::new(ErrorKind::Internal, why)
});
}
let mut ready: Vec<_> = ready.iter_mut().collect();
ready.sort_unstable_by_key(|item| item.effect.id());
for item in ready {
if item.message || item.decision.is_decided() {
continue;
}
let ending = item.of.and_then(|of| turns.get(of.0).ok());
if ending.and_then(TurnOutcome::stops).is_some() {
continue;
}
let (entity, effect) = (item.entity, item.effect);
let changes = item.tool.is_some() && !item.footprint.is_some_and(|it| it.0.reads_only());
if let Some(unsaved) = unsaved.as_ref().filter(|_| changes) {
commands
.entity(entity)
.try_insert_if_new(Outcome(Err(unsaved.clone())));
continue;
}
let agent = item.agent.map(|agent| agent.0);
let caller = agent.and_then(|agent| agents.get(agent).ok());
let scope = caller.map(|id| id.0.as_str()).unwrap_or_default();
let origin = DispatchContext {
parent: effect.parent,
scope: caller.map(|id| Arc::from(id.0.as_str())),
purpose: effect.purpose.clone(),
};
let mut started = commands.entity(entity);
if let Some(mut model) = item.model {
let (connection, mut dropped) = match models.start(&mut model, item.adapters, agent) {
Ok(started) => started,
Err(report) => {
started.try_insert_if_new(Outcome(Err(report)));
continue;
}
};
let kind = EffectKind::Completion {
request: model.request.clone(),
stream: true,
};
dropped.extend(connection.handler.0.check(&kind));
let reference = connection.spec.reference();
let span =
info_span!("model_call", agent = scope, effect = %effect.id, model = reference);
let (sender, feed) = feed();
let wake_up = wake.clone();
let consume = move |reply| stream_reply(reply, sender, wake_up);
let handler = connection.handler.0.clone();
let work = effects.run(effect.id, origin, handler, kind, consume);
let running = Running::finishing(model_pool(), &wake, work.instrument(span), answer);
started.try_insert((
Started { at },
running,
feed,
SentTo(connection),
DroppedOptions(dropped),
));
continue;
}
let (handler, kind) = match (item.request, item.tool) {
(
Some(EffectRequest {
kind,
handler: Some(handler),
}),
_,
) => (handler.clone(), kind.clone()),
(_, Some(ToolRequest { call })) => match tools.start(call, item.handler, agent) {
ToolStart::Run(handler, kind) => (handler, kind),
ToolStart::Open(open, tool) => {
commands.queue(move |world: &mut World| {
if let Err(report) = open(world, tool, entity, Started { at })
&& let Ok(mut call) = world.get_entity_mut(entity)
{
call.insert_if_new(Outcome(Err(report)));
}
});
continue;
}
ToolStart::Refused(report) => {
started.try_insert_if_new(Outcome(Err(report)));
continue;
}
},
_ => {
started.try_insert(Started { at });
continue;
}
};
let span = match &kind {
EffectKind::ToolCall { name, .. } => {
info_span!("tool_call", agent = scope, tool = name.as_str(), parent = ?effect.parent)
}
other => info_span!("effect", kind = other.name()),
};
let work = effects.run(effect.id, origin, handler.0, kind, Reply::into_outcome);
let running = Running::finishing(effect_pool(), &wake, work.instrument(span), answer);
started.try_insert((Started { at }, running));
}
}
fn answer(mut entity: EntityWorldMut, (outcome, published): Ran) {
if entity.contains::<Outcome>() {
return;
}
if let Some(published) = published {
entity.insert(Published(published));
}
entity.insert(Outcome(outcome));
}
#[derive(QueryData)]
#[query_data(mutable)]
struct StartItem {
entity: Entity,
effect: &'static Effect,
decision: &'static Decision,
request: Option<&'static EffectRequest>,
tool: Option<&'static ToolRequest>,
model: Option<&'static mut ModelRequest>,
adapters: Option<&'static Adapters>,
turn_call: Has<TurnCall>,
message: Has<Incoming>,
handler: Option<&'static CallHandler>,
footprint: Option<&'static Footprint>,
agent: Option<&'static ForAgent>,
of: Option<&'static CallOf>,
}
fn effect_pool() -> &'static AsyncComputeTaskPool {
AsyncComputeTaskPool::get_or_init(TaskPool::default)
}
type Due = (
With<Effect>,
Without<Settled>,
Or<((With<Outcome>, Without<Holds>), With<SupersededBy>)>,
);
fn settle(
world: &mut World,
due: &mut QueryState<(Entity, &Effect), Due>,
live: &mut QueryState<&Effect, Without<Settled>>,
) {
let mut settling: Vec<(EffectId, Entity)> = due
.iter(world)
.map(|(entity, effect)| (effect.id, entity))
.collect();
if let Some(failure) = failed_in(Stage::Review) {
let what = "nothing settled this pass";
let notice = world
.get_resource_mut::<Passes>()
.filter(|_| !settling.is_empty())
.and_then(|mut passes| closed(&mut passes, failure, what));
if let Some(notice) = notice {
world.write_message(notice);
}
return;
}
settling.sort_unstable();
for (_, entity) in settling {
let superseded = world.get::<SupersededBy>(entity).is_some();
settle_now(world, entity, superseded.then_some(Disposition::Superseded));
}
record_nested(world, live);
}
fn settle_now(world: &mut World, entity: Entity, disposition: Option<Disposition>) {
let record = assemble(world, entity, disposition);
if let Ok(mut effect) = world.get_entity_mut(entity) {
effect.remove::<Running>().insert(Settled);
}
if let Some(record) = record {
world.write_message(EffectRecorded(record));
}
}
pub(crate) fn release(commands: &mut Commands, holds: Option<&Holds>) {
for hold in holds.into_iter().flat_map(Holds::iter) {
commands.entity(hold).try_despawn();
}
}
pub(crate) fn cancel(world: &mut World, entity: Entity, why: &str) {
let unsettled = world
.get_entity(entity)
.is_ok_and(|effect| effect.contains::<Effect>() && !effect.contains::<Settled>());
if !unsettled {
return;
}
if let Ok(mut effect) = world.get_entity_mut(entity) {
effect.despawn_related::<Holds>();
}
let Ok(mut effect) = world.get_entity_mut(entity) else {
return;
};
if !effect.contains::<Outcome>() {
let stopped = match effect.get::<IfStopped>() {
Some(IfStopped(outcome)) => outcome.clone(),
None => Outcome(Err(ErrorReport::new(ErrorKind::Cancelled, why))),
};
effect.insert(stopped);
}
settle_now(world, entity, Some(Disposition::Cancelled));
}
fn record_nested(world: &mut World, live: &mut QueryState<&Effect, Without<Settled>>) {
let Some(buffer) = world
.get_resource::<Effects>()
.map(|effects| effects.buffer().clone())
else {
return;
};
if buffer.resolved() == 0 {
return;
}
let live: HashSet<EffectId> = live.iter(world).map(Effect::id).collect();
for record in buffer.take_resolved_where(|id| !live.contains(&id)) {
world.write_message(EffectRecorded(record));
}
}
fn assemble(
world: &World,
entity: Entity,
disposition: Option<Disposition>,
) -> Option<EffectRecord> {
let buffer = world.get_resource::<Effects>()?.buffer();
let effect = world.get::<Effect>(entity)?;
let current = world
.get::<Outcome>(entity)
.map(|outcome| outcome.0.clone());
let reviewed = world
.get_entity(entity)
.ok()
.and_then(|entity| entity.get_change_ticks::<Outcome>())
.is_some_and(|ticks| ticks.changed != ticks.added);
let (outcome, reviewed) = match (reviewed, &effect.raw, current) {
(true, Some(raw), current) => (Some(raw.clone()), current),
(_, _, current) => (current, None),
};
let disposition = disposition
.or(effect.answered.filter(|_| outcome.is_some()))
.unwrap_or(Disposition::Cancelled);
let outcome = outcome.unwrap_or_else(|| match disposition {
Disposition::Superseded => Err(ErrorReport::new(
ErrorKind::Cancelled,
"superseded before it was answered",
)),
_ => Err(cancelled()),
});
let id = effect.id;
let handled = buffer.take(id).or_else(|| {
buffer.contains(id).then(|| {
buffer.resolve(id, outcome.clone());
buffer.take(id)
})?
});
let mut record = match handled {
Some(mut record) => {
if disposition != Disposition::Ran {
record.outcome = outcome;
}
record
}
None => {
let (key, kind) = kind_of(world, entity)?;
EffectRecord {
published: world
.get::<Published>(entity)
.map(|published| published.0.clone()),
id,
key,
kind,
outcome,
events: None,
stream_origin: None,
parent: effect.parent,
scope: scope_of(world, entity),
purpose: effect.purpose.clone(),
turn: None,
started_ms: world
.get::<Started>(entity)
.map(|started| millis(started.at)),
finished_ms: Some(now_ms()),
reviewed: None,
disposition: Disposition::Ran,
decided_by: None,
}
}
};
record.reviewed = reviewed;
record.disposition = disposition;
record.decided_by.clone_from(&effect.decided_by);
record.turn = world
.get::<CallOf>(entity)
.and_then(|call| world.get::<Turn>(call.0))
.map(|turn| turn.index);
Some(record)
}
fn kind_of(world: &World, entity: Entity) -> Option<(HandlerKey, EffectKind)> {
if let Some(ToolRequest { call }) = world.get::<ToolRequest>(entity) {
let name = call.function.name.as_str();
let key = world
.get::<CallHandler>(entity)
.map_or_else(|| tool_key(name), |handler| handler.0.0.key().clone());
return Some((key, EffectKind::from(call)));
}
if let Some(model) = world.get::<ModelRequest>(entity) {
let key = model.connection.as_ref().map_or_else(
|| HandlerKey::new("completion"),
|connection| connection.handler.0.key().clone(),
);
let kind = EffectKind::Completion {
request: model.request.clone(),
stream: true,
};
return Some((key, kind));
}
if let Some(message) = world.get::<Incoming>(entity) {
return Some(message_kind(message));
}
let request = world.get::<EffectRequest>(entity)?;
let key = match &request.handler {
Some(handler) => handler.0.key().clone(),
None => match &request.kind {
EffectKind::ToolCall { name, .. } => tool_key(name),
other => HandlerKey::new(other.name()),
},
};
Some((key, request.kind.clone()))
}
fn scope_of(world: &World, entity: Entity) -> Option<Arc<str>> {
let agent = world.get::<ForAgent>(entity)?.0;
world
.get::<AgentId>(agent)
.map(|id| Arc::from(id.0.as_str()))
}