1use std::borrow::Cow;
18use std::fmt;
19use std::mem;
20use std::pin::Pin;
21use std::sync::Arc;
22use std::time::{Duration, UNIX_EPOCH};
23
24use bevy_ecs::lifecycle::HookContext;
25use bevy_ecs::prelude::*;
26use bevy_ecs::query::QueryData;
27use bevy_ecs::system::SystemParam;
28use bevy_ecs::world::DeferredWorld;
29use bevy_log::{info, warn};
30use bevy_reflect::prelude::*;
31use bevy_tasks::futures::check_ready;
32use bevy_tasks::{IoTaskPool, TaskPool};
33use bevy_time::DelayedCommandsExt;
34use crossbeam_channel::{Receiver, Sender};
35use futures::StreamExt;
36use rig_core::completion::message::turn_failure;
37use rig_core::completion::{
38 self, CompletionRequest, CompletionResponse, GenerationOptions, Message, Section,
39 UnsupportedOption,
40};
41use rig_core::conversation::pairing::close_pending_with;
42use rig_core::conversation::{Author, Halt, TurnEnd};
43use rig_core::effect::{self, EffectId, EffectKind};
44use rig_core::error::{ErrorKind, ErrorReport};
45use rig_core::message::{AssistantContent, AssistantMessage, ToolCall, UserContent};
46use rig_core::now_ms;
47use rig_core::serve::{Reply, stream_truncated};
48use rig_core::streaming::{Item, Relayed, StreamEvent};
49use rig_core::tool::{ToolErrorKind, ToolExecutionError, ToolResult, feedback};
50use rig_memory::Anchor;
51use web_time::Instant;
52
53use super::agent::{
54 ActiveTurn, Agent, AgentId, CacheKey, CallOf, Calls, Conversation, EffectParent, LastUsage,
55 Notice, Partial, Retry, Spawned, StopTurn, SystemPrompt, ToolAccess, Turn, TurnOf, TurnOutcome,
56};
57use super::inbox::{DeliveryMode, Inboxes, carry_on, deliver};
58use super::model::{Connection, ReasoningSetting};
59use super::prompt::{PromptSection, SectionOf};
60use super::restore::{Restored, commit_results, mark_settled};
61use super::retry::{RetryPolicy, Verdict};
62use super::session_log::{ConversationChanged, SessionLog, save_conversation};
63#[cfg(doc)]
64use super::stage::Stage;
65use super::stage::{
66 Admitted, Effect, ForAgent, HoldOf, Holds, IfStopped, Outcome, Settled, Started, SupersededBy,
67 Waiting, cancel, release,
68};
69use super::tasks::{Running, Wake};
70use super::tools::{Footprint, ToolDef, ToolRequest};
71
72pub(crate) const NO_MODEL: &str = "No model is connected; pick one first.";
74
75pub const RETRY_POLICY: RetryPolicy = RetryPolicy::DEFAULT;
78
79#[derive(Component, Reflect, Clone, Copy, Debug, Default)]
81#[reflect(Component, Default)]
82pub struct Recovery {
83 pub retries: u32,
85}
86
87#[derive(Component, Reflect, Clone, Debug)]
93#[reflect(Component, from_reflect = false)]
94pub struct Backoff {
95 pub attempt: u32,
97 #[reflect(ignore)]
99 pub until: Instant,
100}
101
102impl Backoff {
103 pub fn seconds_left(&self) -> u64 {
105 let left = self.until.saturating_duration_since(Instant::now());
106 left.as_secs() + u64::from(left.subsec_nanos() > 0)
107 }
108}
109
110#[derive(Component, Reflect, Clone)]
124#[reflect(opaque, Component, Clone)]
125#[require(Effect, Adapters)]
126pub struct ModelRequest {
127 pub request: Arc<CompletionRequest>,
130 pub connection: Option<Connection>,
133 pub prompt: Option<completion::SystemPrompt>,
137 pub history: Option<rig_core::conversation::Conversation>,
150}
151
152impl ModelRequest {
153 pub fn new(request: impl Into<Arc<CompletionRequest>>) -> Self {
155 Self {
156 request: request.into(),
157 connection: None,
158 prompt: None,
159 history: None,
160 }
161 }
162
163 pub fn request_mut(&mut self) -> &mut CompletionRequest {
166 Arc::make_mut(&mut self.request)
167 }
168
169 #[must_use]
171 pub fn on(self, connection: Connection) -> Self {
172 Self {
173 connection: Some(connection),
174 ..self
175 }
176 }
177}
178
179#[derive(Component, Clone)]
188pub struct SentTo(pub Connection);
189
190#[derive(Component, Clone, Debug, Default)]
196pub struct DroppedOptions(pub Vec<UnsupportedOption>);
197
198#[derive(Component, Clone, Default, Debug)]
204pub struct Adapters(pub Vec<Adapter>);
205
206impl Adapters {
207 pub fn apply(&self, connection: Connection) -> Connection {
209 let placed = |outermost| self.0.iter().filter(move |it| it.outermost == outermost);
210 placed(false)
211 .chain(placed(true))
212 .fold(connection, |connection, adapter| adapter.apply(connection))
213 }
214}
215
216#[derive(Clone)]
223pub struct Adapter {
224 pub name: Cow<'static, str>,
226 outermost: bool,
227 apply: Arc<dyn Fn(Connection) -> Connection + Send + Sync>,
228}
229
230impl Adapter {
231 pub fn new(
233 name: impl Into<Cow<'static, str>>,
234 apply: impl Fn(Connection) -> Connection + Send + Sync + 'static,
235 ) -> Self {
236 Self {
237 name: name.into(),
238 outermost: false,
239 apply: Arc::new(apply),
240 }
241 }
242
243 #[must_use]
247 pub fn outermost(self) -> Self {
248 Self {
249 outermost: true,
250 ..self
251 }
252 }
253
254 pub fn apply(&self, connection: Connection) -> Connection {
256 (self.apply)(connection)
257 }
258}
259
260impl fmt::Debug for Adapter {
261 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
262 f.debug_tuple("Adapter").field(&self.name).finish()
263 }
264}
265
266#[derive(Component, Reflect, Clone, Copy, Debug, Default)]
272#[reflect(Component, Default)]
273#[require(Effect, Partial, Name = Name::new("model call"))]
274pub struct TurnCall;
275
276#[derive(Component, Reflect, Clone, Copy, Debug)]
282#[reflect(Component)]
283#[require(Effect)]
284pub struct SideCall(pub Entity);
285
286#[derive(Component, Clone, Debug)]
294pub struct RetryVerdict(pub Verdict);
295
296#[derive(Component)]
298pub(crate) struct Feed(Receiver<Fed>);
299
300pub(crate) struct Fed(web_time::SystemTime, Option<StreamEvent>);
303
304#[derive(Message, Clone, Debug)]
313pub struct StreamDelta {
314 pub effect: Entity,
316 pub agent: Option<Entity>,
318 pub event: StreamEvent,
320}
321
322#[derive(Component, Clone, Copy, Debug)]
328pub struct StreamTiming {
329 pub first: web_time::SystemTime,
331 pub last: web_time::SystemTime,
334 pub items: u64,
336}
337
338pub(crate) fn model_pool() -> &'static IoTaskPool {
340 IoTaskPool::get_or_init(TaskPool::default)
341}
342
343pub(crate) fn spawn_turn_call(commands: &mut Commands, turn: Entity, agent: Entity) {
345 commands.spawn((TurnCall, CallOf(turn), ForAgent(agent)));
346}
347
348pub(crate) fn on_retry(
351 retry: On<Retry>,
352 mut agents: Query<(&mut Conversation, Has<ActiveTurn>), With<Agent>>,
353 mut commands: Commands,
354 mut notices: MessageWriter<Notice>,
355) {
356 let agent = retry.entity;
357 let Ok((mut conversation, busy)) = agents.get_mut(agent) else {
358 return;
359 };
360 if busy {
361 notices.write(Notice::turn_running(agent));
362 return;
363 }
364 if !conversation.resume() {
365 notices.write(Notice::info(
366 agent,
367 "Nothing to retry: the model answered the last message.",
368 ));
369 return;
370 }
371 let turn = commands.spawn(TurnOf(agent)).id();
372 spawn_turn_call(&mut commands, turn, agent);
373}
374
375pub(crate) fn start_turn(mut world: DeferredWorld, context: HookContext) {
381 let Some(&TurnOf(agent)) = world.get::<TurnOf>(context.entity) else {
382 return;
383 };
384 world
385 .commands()
386 .entity(context.entity)
387 .try_insert(ForAgent(agent));
388 let restored = world.get::<Restored>(context.entity).is_some();
389 let Some(mut conversation) = world.get_mut::<Conversation>(agent) else {
390 return;
391 };
392 let last = conversation.turns().next_back().map(|(_, record)| record);
393 let record = match last.filter(|record| restored && record.end.is_none()) {
394 Some(record) => record,
395 None => conversation.open_turn(),
396 };
397 if let Some(mut turn) = world.get_mut::<Turn>(context.entity) {
398 turn.index = record.index;
399 turn.start = record.start;
400 }
401}
402
403fn turn_end(outcome: &TurnOutcome, conversation: &Conversation) -> TurnEnd {
406 match outcome {
407 TurnOutcome::Answered(_) => {
408 let last = conversation.entries().next_back();
409 let answer = last.filter(|entry| matches!(entry.message, Message::Assistant(_)));
410 TurnEnd::Answered {
411 answer: answer.map(|entry| entry.id),
412 }
413 }
414 TurnOutcome::Failed(why) => TurnEnd::Failed { why: why.clone() },
415 TurnOutcome::Stopped(why) => TurnEnd::Stopped { why: why.clone() },
416 }
417}
418
419#[derive(Component, Clone, Debug)]
422pub(crate) struct Ending {
423 notice: String,
425 take_out: bool,
427}
428
429const STOPPED_BY_USER: &str = "the user stopped the turn";
431
432pub(crate) fn on_stop_turn(
448 interrupt: On<StopTurn>,
449 spawned: Query<&Spawned>,
450 mut agents: Query<(&AgentId, &mut Conversation, &ActiveTurn)>,
451 mut turns: Query<(&mut Turn, Option<&Calls>)>,
452 replies: Query<ReplyCall>,
453 partials: Query<&Partial, (With<TurnCall>, Without<Settled>, Without<SupersededBy>)>,
454 log: Res<SessionLog>,
455 mut committed: MessageWriter<ConversationChanged>,
456 mut commands: Commands,
457 mut notices: MessageWriter<Notice>,
458) {
459 let root = interrupt.entity;
460 let reason = interrupt.reason.clone();
461 let why = reason.clone().unwrap_or_else(|| STOPPED_BY_USER.to_owned());
462 let below = spawned.iter_descendants::<Spawned>(root);
463 let stopping: Vec<Entity> = std::iter::once(root).chain(below).collect();
464 let mut stopped = false;
465 for agent in stopping {
466 let Ok((id, mut conversation, active)) = agents.get_mut(agent) else {
467 continue;
468 };
469 let turn = active.turn();
470 let (reply, calls) = match turns.get_mut(turn) {
471 Ok((mut state, calls)) => (state.reply.take(), calls),
472 Err(_) => (None, None),
473 };
474 let streamed = partials.iter_many(calls.into_iter().flat_map(Calls::iter));
475 let (mut reasoning, mut text) = (String::new(), String::new());
476 for partial in streamed.flatten() {
477 reasoning.push_str(&partial.reasoning);
478 text.push_str(&partial.text);
479 }
480 let reasoning =
481 (!reasoning.trim().is_empty()).then(|| AssistantContent::reasoning(reasoning));
482 let text = (!text.trim().is_empty()).then(|| AssistantContent::text(text));
483 let content: Vec<_> = reasoning.into_iter().chain(text).collect();
484 if !content.is_empty() {
485 let aborted = AssistantMessage::aborted(None, content, format!("stopped: {why}"));
486 conversation.push(Message::Assistant(aborted));
487 }
488 let reply = reply
489 .zip(calls)
490 .map(|(reply, calls)| reply_calls(reply, calls, &replies));
491 let stopped_call =
492 |call: &ReplyCallItem| Some(not_finished(&why, call.started, call.footprint));
493 if let Some(results) = reply_results(reply.unwrap_or_default(), stopped_call) {
494 commit_results(&mut conversation, results);
495 }
496 conversation.halt(Halt::Stopped);
497 conversation.end_turn(TurnEnd::Stopped {
498 why: reason.clone(),
499 });
500 let conversation = conversation.bypass_change_detection();
501 committed.write_batch(save_conversation(&log, agent, id, conversation));
502 let calls: Vec<Entity> = calls.into_iter().flat_map(Calls::iter).collect();
503 let (why, reason) = (why.clone(), reason.clone());
504 commands.queue(move |world: &mut World| {
505 if let Ok(mut turn) = world.get_entity_mut(turn) {
508 turn.insert(TurnOutcome::Stopped(reason));
509 }
510 for call in calls {
511 let said = world.get_entity(call).ok().and_then(|call| {
512 let footprint = call.get::<Footprint>();
513 call.contains::<ToolRequest>()
514 .then(|| not_finished_text(&why, call.contains::<Started>(), footprint))
515 });
516 cancel(world, call, said.as_deref().unwrap_or(&why));
517 }
518 if let Ok(mut turn) = world.get_entity_mut(turn) {
519 turn.despawn_related::<Holds>();
520 }
521 if let Ok(mut turn) = world.get_entity_mut(turn) {
522 turn.remove::<(TurnOf, Ending)>().insert(Settled);
523 }
524 });
525 notices.write(Notice::info(agent, "Stopped."));
526 stopped = true;
527 }
528 if stopped {
531 let _ = log.flush();
532 }
533}
534
535fn not_finished_text(why: &str, started: bool, footprint: Option<&Footprint>) -> String {
539 let footprint = footprint.cloned().unwrap_or_default();
540 feedback::not_finished(why, started, &footprint.0)
541}
542
543pub(crate) fn not_finished(why: &str, started: bool, footprint: Option<&Footprint>) -> ToolResult {
546 let text = not_finished_text(why, started, footprint);
547 ToolResult::failed(ToolExecutionError::other(text))
548}
549
550fn reply_results(
555 calls: Vec<ReplyCallItem>,
556 missing: impl Fn(&ReplyCallItem) -> Option<ToolResult>,
557) -> Option<Message> {
558 let results: Vec<_> = calls
559 .into_iter()
560 .filter_map(|call| {
561 let outcome = call.outcome.or(call.if_stopped.map(|stopped| &stopped.0));
562 let result = outcome
563 .map(Outcome::tool_result)
564 .or_else(|| missing(&call))?;
565 Some(call.request.call.answer(&result))
566 })
567 .collect();
568 (!results.is_empty()).then(|| Message::tool_results(results))
569}
570
571const EXIT_GRACE: Duration = Duration::from_secs(1);
576
577type Cancelling = Vec<Pin<Box<dyn Future<Output = ()>>>>;
578
579#[derive(Resource, Clone, Copy, Debug, Default)]
583pub(crate) struct Exiting;
584
585pub(crate) fn stop_turns_on_exit(world: &mut World) {
590 world.insert_resource(Exiting);
591 let mut cancelling = Cancelling::new();
592 take_running(world, &mut cancelling);
593 let deadline = Instant::now() + EXIT_GRACE;
594 loop {
595 cancelling.retain_mut(|cancel| check_ready(cancel).is_none());
596 if cancelling.is_empty() || cfg!(target_family = "wasm") || Instant::now() >= deadline {
598 break;
599 }
600 #[cfg(not(target_family = "wasm"))]
601 std::thread::sleep(Duration::from_millis(2));
602 }
603 if let Err(error) = world.run_system_cached(log_answered_results) {
604 warn!("could not log the results of stopped turns: {error}");
605 }
606 let turns: Vec<Entity> = world
607 .query_filtered::<Entity, With<TurnOf>>()
608 .iter(world)
609 .collect();
610 for turn in turns {
611 world.despawn(turn);
612 }
613 world.flush();
614}
615
616fn log_answered_results(
619 turns: Query<(&TurnOf, &Turn, &Calls)>,
620 replies: Query<ReplyCall>,
621 mut agents: Query<&mut Conversation>,
622) {
623 for (&TurnOf(agent), state, calls) in &turns {
624 let Some(reply) = state.reply else {
625 continue;
626 };
627 let results = reply_results(reply_calls(reply, calls, &replies), |_| None);
628 if let Ok(mut conversation) = agents.get_mut(agent)
629 && let Some(results) = results
630 {
631 commit_results(&mut conversation, results);
632 }
633 }
634}
635
636fn take_running(world: &mut World, cancelling: &mut Cancelling) {
639 let calls: Vec<Entity> = world
640 .query_filtered::<Entity, With<Running>>()
641 .iter(world)
642 .collect();
643 for call in calls {
644 if let Ok(mut call) = world.get_entity_mut(call)
645 && let Some(running) = call.take::<Running>()
646 {
647 cancelling.push(Box::pin(running.cancel()));
648 }
649 }
650}
651
652#[derive(QueryData)]
654#[query_data(mutable)]
655pub(crate) struct Builder {
656 id: &'static AgentId,
657 cache: Option<&'static CacheKey>,
658 connection: Option<&'static Connection>,
659 effort: &'static ReasoningSetting,
660 prompt: &'static SystemPrompt,
661 access: &'static ToolAccess,
662 conversation: &'static mut Conversation,
663 parent: Option<&'static EffectParent>,
664}
665
666#[derive(QueryData)]
668#[query_data(mutable)]
669pub(crate) struct Unbuilt {
670 entity: Entity,
671 effect: &'static mut Effect,
672 from: Option<&'static SideCall>,
673 turn_call: Has<TurnCall>,
674 of: Option<&'static CallOf>,
675 agent: Option<&'static ForAgent>,
676}
677
678#[derive(SystemParam)]
681pub(crate) struct Parts<'w, 's> {
682 tools: Query<'w, 's, &'static ToolDef>,
683 sections: Query<'w, 's, (&'static PromptSection, Option<&'static SectionOf>)>,
684}
685
686pub(crate) fn build_requests(
694 mut unbuilt: Query<
695 Unbuilt,
696 (
697 Or<(With<TurnCall>, With<SideCall>)>,
698 Added<Admitted>,
699 Without<ModelRequest>,
700 Without<Outcome>,
701 Without<Started>,
702 ),
703 >,
704 turns: Query<&TurnOf>,
705 mut agents: Query<Builder>,
706 inboxes: Inboxes,
707 parts: Parts,
708 mut commands: Commands,
709 mut notices: MessageWriter<Notice>,
710) {
711 for mut call in &mut unbuilt {
712 let turn = call.of.map(|of| of.0);
713 let agent = call
714 .from
715 .map(|from| from.0)
716 .or(call.agent.map(|agent| agent.0))
717 .or_else(|| turns.get(turn?).ok().map(|of| of.0));
718 let mut built = commands.entity(call.entity);
719 let Some((agent, mut builder)) =
720 agent.and_then(|agent| Some((agent, agents.get_mut(agent).ok()?)))
721 else {
722 let why = ErrorReport::new(ErrorKind::Internal, "a model call needs an agent");
723 built.try_insert_if_new(Outcome(Err(why)));
724 continue;
725 };
726 if call.agent.is_none() {
727 built.try_insert(ForAgent(agent));
728 }
729 let effect = &mut *call.effect;
730 if effect.parent.is_none() {
731 effect.parent = builder.parent.map(|parent| parent.0);
732 }
733 let spec = builder.connection.map(|connection| &*connection.spec);
734 if call.turn_call {
735 effect.purpose.get_or_insert("turn".into());
736 let waiting = inboxes.of(agent).into_iter();
737 let waiting = waiting.filter(|(_, message)| message.mode != DeliveryMode::Queue);
738 let conversation = &mut *builder.conversation;
739 let into = (turn, &mut commands);
740 deliver(agent, spec, waiting, conversation, &mut notices, into);
741 }
742 if call.turn_call && builder.conversation.messages().is_empty() {
743 let why = ErrorReport::new(ErrorKind::Internal, "The conversation is empty.");
744 commands
745 .entity(call.entity)
746 .try_insert_if_new(Outcome(Err(why)));
747 continue;
748 }
749 let history = builder.conversation.request_view();
750 let model = build((agent, &builder), spec, &parts, history);
751 commands.entity(call.entity).try_insert(model);
752 }
753}
754
755fn build(
758 (agent, builder): (Entity, &BuilderItem),
759 spec: Option<&rig_core::catalog::ModelSpec>,
760 parts: &Parts,
761 history: rig_core::conversation::Conversation,
762) -> ModelRequest {
763 let mut offered: Vec<&ToolDef> = parts
766 .tools
767 .iter()
768 .filter(|def| builder.access.allows(def.0.name.as_str()))
769 .collect();
770 offered.sort_by(|a, b| a.0.name.as_str().cmp(b.0.name.as_str()));
771 let mut sections: Vec<&PromptSection> = parts
773 .sections
774 .iter()
775 .filter(|(_, of)| of.is_none_or(|of| of.0 == agent))
776 .map(|(section, _)| section)
777 .collect();
778 sections.sort_by(|a, b| (a.order, &a.tag).cmp(&(b.order, &b.tag)));
779 let sections = sections
780 .into_iter()
781 .map(|section| Section::new(§ion.tag, §ion.text));
782 let prompt = completion::SystemPrompt {
783 role: builder.prompt.0.clone(),
784 sections: sections.collect(),
785 };
786 let effort = builder.effort.0;
787 let options = spec.map_or_else(
788 || {
789 let options = GenerationOptions::default();
790 match effort {
791 Some(reasoning) => options.reasoning(reasoning),
792 None => options,
793 }
794 },
795 |spec| spec.default_options(effort),
796 );
797 let request = CompletionRequest::from(Vec::new())
798 .tools(offered.iter().map(|def| def.0.clone()))
799 .options(options.cache_key(CacheKey::of(builder.cache, builder.id)));
800 ModelRequest {
801 request: request.into(),
802 connection: builder.connection.cloned(),
803 prompt: Some(prompt),
804 history: Some(history),
805 }
806}
807
808#[derive(SystemParam)]
812pub struct ModelStarts<'w, 's> {
813 agents: Query<'w, 's, &'static Connection, With<Agent>>,
814}
815
816impl ModelStarts<'_, '_> {
817 pub fn sent(
824 &self,
825 model: &ModelRequest,
826 adapters: Option<&Adapters>,
827 agent: Option<Entity>,
828 ) -> Result<(Connection, CompletionRequest), ErrorReport> {
829 let (connection, request, _) = self.sending(model, adapters, agent)?;
830 Ok((connection, request))
831 }
832
833 pub fn dropped(
839 &self,
840 model: &ModelRequest,
841 adapters: Option<&Adapters>,
842 agent: Option<Entity>,
843 ) -> Result<Vec<UnsupportedOption>, ErrorReport> {
844 let (connection, request, mut dropped) = self.sending(model, adapters, agent)?;
845 let kind = EffectKind::Completion {
846 request: request.into(),
847 stream: true,
848 };
849 dropped.extend(connection.handler.0.check(&kind));
850 Ok(dropped)
851 }
852
853 pub fn prompt(
859 &self,
860 model: &ModelRequest,
861 adapters: Option<&Adapters>,
862 agent: Option<Entity>,
863 ) -> Result<(Connection, Option<completion::SystemPrompt>), ErrorReport> {
864 let connection = self.connection(model.connection.as_ref(), adapters, agent)?;
865 Ok((connection, model.prompt.clone()))
866 }
867
868 fn sending(
871 &self,
872 model: &ModelRequest,
873 adapters: Option<&Adapters>,
874 agent: Option<Entity>,
875 ) -> Result<(Connection, CompletionRequest, Vec<UnsupportedOption>), ErrorReport> {
876 let mut model = model.clone();
877 let (connection, dropped) = self.start(&mut model, adapters, agent)?;
878 Ok((connection, Arc::unwrap_or_clone(model.request), dropped))
879 }
880
881 pub(crate) fn start(
884 &self,
885 model: &mut ModelRequest,
886 adapters: Option<&Adapters>,
887 agent: Option<Entity>,
888 ) -> Result<(Connection, Vec<UnsupportedOption>), ErrorReport> {
889 let request = Arc::make_mut(&mut model.request);
890 let to = model.connection.as_ref();
891 let started = self.fit(request, to, adapters, agent)?;
892 let history = match model.history.take() {
893 Some(history) => history.into_messages(),
894 None => mem::take(&mut request.chat_history),
895 };
896 let prompt = model.prompt.take();
898 let prompt = prompt.map(|prompt| Arc::new(Message::system(prompt.text())));
899 request.chat_history = prompt.into_iter().chain(history).collect();
900 Ok(started)
901 }
902
903 fn connection(
906 &self,
907 to: Option<&Connection>,
908 adapters: Option<&Adapters>,
909 agent: Option<Entity>,
910 ) -> Result<Connection, ErrorReport> {
911 let own = agent.and_then(|agent| self.agents.get(agent).ok());
912 let Some(connection) = to.or(own).cloned() else {
913 return Err(ErrorReport::new(ErrorKind::HandlerUnavailable, NO_MODEL));
914 };
915 Ok(match adapters {
916 Some(adapters) => adapters.apply(connection),
917 None => connection,
918 })
919 }
920
921 fn fit(
926 &self,
927 request: &mut CompletionRequest,
928 to: Option<&Connection>,
929 adapters: Option<&Adapters>,
930 agent: Option<Entity>,
931 ) -> Result<(Connection, Vec<UnsupportedOption>), ErrorReport> {
932 let connection = self.connection(to, adapters, agent)?;
933 let spec = &connection.spec;
934 if !spec.tools {
935 request.tools.clear();
936 }
937 if request.options.cache.is_none() {
941 request.options.cache = spec.default_options(request.options.reasoning).cache;
942 }
943 let dropped = spec.fit(request);
944 Ok((connection, dropped))
945 }
946}
947
948pub(crate) async fn stream_reply(
952 reply: Reply,
953 feed: Sender<Fed>,
954 wake: Wake,
955) -> Result<effect::Answered, ErrorReport> {
956 let mut events = reply.into_stream();
957 while let Some(item) = events.next().await {
958 let event = match item {
959 Err(report) => return Err(report),
960 Ok(Relayed::Origin(_)) => continue,
962 Ok(Relayed::Done(response)) => {
963 feed.send(Fed(web_time::SystemTime::now(), None)).ok();
965 return Ok(effect::Answered::Completion(response));
966 }
967 Ok(Relayed::Item(Item::Event(event))) => Some(event),
968 Ok(Relayed::Item(_)) => None,
969 };
970 feed.send(Fed(web_time::SystemTime::now(), event)).ok();
971 wake.wake();
972 }
973 Err(stream_truncated())
974}
975
976pub(crate) fn feed() -> (Sender<Fed>, Feed) {
979 let (sender, receiver) = crossbeam_channel::unbounded();
980 (sender, Feed(receiver))
981}
982
983pub(crate) fn stream_partials(
988 mut calls: Query<(
989 Entity,
990 &Feed,
991 Option<&mut Partial>,
992 Option<&mut StreamTiming>,
993 Option<&ForAgent>,
994 )>,
995 mut deltas: MessageWriter<StreamDelta>,
996 mut commands: Commands,
997) {
998 for (effect, feed, mut partial, mut streamed, agent) in &mut calls {
999 let mut first: Option<StreamTiming> = None;
1001 for Fed(at, event) in feed.0.try_iter() {
1002 match streamed.as_deref_mut().or(first.as_mut()) {
1003 Some(stamp) => {
1004 stamp.last = at;
1005 stamp.items += 1;
1006 }
1007 None => {
1008 first = Some(StreamTiming {
1009 first: at,
1010 last: at,
1011 items: 1,
1012 });
1013 }
1014 }
1015 let Some(event) = event else {
1016 continue;
1017 };
1018 if let Some(partial) = partial.as_mut() {
1019 match &event {
1020 StreamEvent::Text { text, .. } => partial.text.push_str(text),
1021 StreamEvent::Reasoning { text, .. } => partial.reasoning.push_str(text),
1022 _ => {}
1023 }
1024 }
1025 let agent = agent.map(|agent| agent.0);
1026 deltas.write(StreamDelta {
1027 effect,
1028 agent,
1029 event,
1030 });
1031 }
1032 if let Some(first) = first {
1033 commands.entity(effect).try_insert(first);
1034 }
1035 }
1036}
1037
1038pub(crate) fn judge_failures(
1042 failed: Query<
1043 (Entity, &Effect, &Outcome, Option<&CallOf>),
1044 (With<ModelRequest>, Added<Outcome>, Without<Settled>),
1045 >,
1046 turns: Query<&Recovery>,
1047 mut commands: Commands,
1048) {
1049 for (call, effect, Outcome(outcome), of) in &failed {
1050 let Err(report) = outcome else {
1051 continue;
1052 };
1053 if !effect.ran() {
1054 continue;
1055 }
1056 let retries = of
1057 .and_then(|of| turns.get(of.0).ok())
1058 .map_or(0, |recovery| recovery.retries);
1059 let now = UNIX_EPOCH + Duration::from_millis(now_ms());
1062 let verdict = RETRY_POLICY.verdict(report, retries, now);
1063 commands.entity(call).try_insert(RetryVerdict(verdict));
1064 }
1065}
1066
1067pub(crate) fn back_off(
1073 failed: Query<
1074 (Entity, &Effect, &Outcome, &RetryVerdict, &CallOf),
1075 (
1076 With<TurnCall>,
1077 Without<Holds>,
1078 Without<Settled>,
1079 Without<SupersededBy>,
1080 ),
1081 >,
1082 mut turns: Query<(&TurnOf, &mut Recovery, Option<&TurnOutcome>)>,
1083 mut commands: Commands,
1084 mut notices: MessageWriter<Notice>,
1085) {
1086 for (call, effect, Outcome(outcome), RetryVerdict(verdict), &CallOf(turn)) in &failed {
1087 let (Err(report), Verdict::Retry(delay)) = (outcome, verdict) else {
1088 continue;
1089 };
1090 let Ok((&TurnOf(agent), mut recovery, ending)) = turns.get_mut(turn) else {
1091 continue;
1092 };
1093 if ending.and_then(TurnOutcome::stops).is_some() {
1094 continue;
1095 }
1096 recovery.retries += 1;
1097 let backoff = Backoff {
1098 attempt: recovery.retries,
1099 until: Instant::now() + *delay,
1100 };
1101 notices.write(Notice::info(
1102 agent,
1103 format!(
1104 "The model call failed: {report}. Retrying in {}s ({}/{}).",
1105 backoff.seconds_left(),
1106 backoff.attempt,
1107 RETRY_POLICY.max_retries
1108 ),
1109 ));
1110 let hold = commands
1111 .spawn((Name::new("retry wait"), backoff, HoldOf(call)))
1112 .id();
1113 let parent = effect.parent;
1114 commands
1115 .delayed()
1116 .duration(*delay)
1117 .queue(move |world: &mut World| {
1118 if world.get_entity(hold).is_err() {
1120 return;
1121 }
1122 world.despawn(hold);
1123 world
1124 .spawn((
1125 TurnCall,
1126 CallOf(turn),
1127 ForAgent(agent),
1128 Effect::with_parent_id(parent),
1129 ))
1130 .add_related::<SupersededBy>(&[call]);
1131 });
1132 }
1133}
1134
1135pub(crate) fn stop_calls(
1142 turns: Query<(&TurnOutcome, &Calls), With<TurnOf>>,
1143 calls: Query<
1144 (
1145 Has<Outcome>,
1146 Option<&Holds>,
1147 Option<&IfStopped>,
1148 (Has<ToolRequest>, Has<Started>, Option<&Footprint>),
1149 ),
1150 (With<Admitted>, Without<Settled>),
1151 >,
1152 mut commands: Commands,
1153) {
1154 for (ending, of) in &turns {
1155 let Some(why) = ending.stops() else {
1156 continue;
1157 };
1158 for (call, (answered, holds, if_stopped, (tool, started, footprint))) in of
1159 .iter()
1160 .filter_map(|call| Some((call, calls.get(call).ok()?)))
1161 {
1162 release(&mut commands, holds);
1163 if !answered {
1164 let said = match tool {
1165 true => not_finished_text(&why, started, footprint),
1166 false => why.clone(),
1167 };
1168 let stopped = match if_stopped {
1169 Some(IfStopped(outcome)) => outcome.clone(),
1170 None => Outcome(Err(ErrorReport::new(ErrorKind::Cancelled, said))),
1171 };
1172 commands.entity(call).try_insert_if_new(stopped);
1173 }
1174 }
1175 }
1176}
1177
1178pub(crate) fn commit_turn_calls(
1189 calls: Query<
1190 (
1191 Entity,
1192 &Effect,
1193 Option<&Outcome>,
1194 Option<&RetryVerdict>,
1195 &CallOf,
1196 Has<SupersededBy>,
1197 ),
1198 (With<TurnCall>, Added<Settled>),
1199 >,
1200 mut turns: Query<(&TurnOf, &mut Recovery, &mut Turn, Option<&TurnOutcome>)>,
1201 mut agents: Query<Answered>,
1202 inboxes: Inboxes,
1203 mut commands: Commands,
1204 mut notices: MessageWriter<Notice>,
1205) {
1206 for (call, effect, outcome, verdict, &CallOf(turn), superseded) in &calls {
1207 commands.entity(call).try_remove::<(Partial, Feed)>();
1209 let Some(Outcome(outcome)) = outcome.filter(|_| !superseded) else {
1210 continue;
1211 };
1212 let Ok((&TurnOf(agent), mut recovery, mut state, ending)) = turns.get_mut(turn) else {
1213 continue;
1214 };
1215 let Ok(mut answered) = agents.get_mut(agent) else {
1216 continue;
1217 };
1218 let stopped = ending.and_then(TurnOutcome::stops);
1219 let propose = |commands: &mut Commands, outcome: TurnOutcome, kernel: Option<Ending>| {
1220 if ending.is_none() {
1221 let mut turn = commands.entity(turn);
1222 turn.try_insert(outcome);
1223 if let Some(kernel) = kernel {
1224 turn.try_insert(kernel);
1225 }
1226 }
1227 };
1228 let report = match outcome {
1229 Ok(effect::Answered::Completion(response)) => {
1230 recovery.retries = 0;
1231 let inbox = (&inboxes, &mut commands);
1232 let next =
1233 commit_reply((agent, turn), response, &mut answered, inbox, &mut notices);
1234 match (next, stopped) {
1235 (Next::End(outcome, kernel), _) => propose(&mut commands, outcome, kernel),
1236 (Next::Call, None) => spawn_turn_call(&mut commands, turn, agent),
1237 (Next::Run(calls), None) => {
1238 state.reply = Some(effect.id());
1239 let cut_off = cut_off(response);
1240 spawn_reply(&mut commands, (agent, turn), calls, effect.id(), &cut_off);
1241 }
1242 (Next::Run(calls), Some(why)) => {
1243 let results = close_pending_with(&calls, &format!("not run: {why}"));
1244 commit_results(&mut answered.conversation, results);
1245 }
1246 (Next::Call, Some(_)) => {}
1247 }
1248 continue;
1249 }
1250 _ if stopped.is_some() => continue,
1251 Ok(other) => ErrorReport::new(
1252 ErrorKind::Internal,
1253 format!("the model answered with a {} outcome", other.family()),
1254 ),
1255 Err(report) => report.clone(),
1256 };
1257 let kept = |notice: String| {
1258 Some(Ending {
1259 notice,
1260 take_out: false,
1261 })
1262 };
1263 let mut propose = |outcome, kernel| propose(&mut commands, outcome, kernel);
1264 match verdict.map(|verdict| &verdict.0) {
1265 None => propose(
1266 TurnOutcome::Failed(report.message.clone()),
1267 kept(report.message),
1268 ),
1269 Some(Verdict::GaveUp(why)) => propose(
1270 TurnOutcome::Failed(report.to_string()),
1271 kept(format!(
1272 "The model call failed: {report}. Not retrying: {why}. Your message is kept \
1273 for a retry."
1274 )),
1275 ),
1276 Some(Verdict::Retry(_)) => propose(
1277 TurnOutcome::Failed(report.to_string()),
1278 kept(format!(
1279 "The model call failed: {report}. Your message is kept for a retry."
1280 )),
1281 ),
1282 Some(verdict @ (Verdict::Overflow | Verdict::Final)) => {
1283 let why = match verdict {
1284 Verdict::Overflow => format!(
1285 "The conversation does not fit the model's context window: {report}"
1286 ),
1287 _ => format!("The model call failed: {report}"),
1288 };
1289 let ending = Ending {
1290 notice: why.clone(),
1291 take_out: true,
1292 };
1293 propose(TurnOutcome::Failed(why), Some(ending));
1294 }
1295 }
1296 }
1297}
1298
1299enum Next {
1301 End(TurnOutcome, Option<Ending>),
1303 Call,
1305 Run(Vec<ToolCall>),
1307}
1308
1309fn commit_reply(
1312 (agent, turn): (Entity, Entity),
1313 response: &CompletionResponse,
1314 answered: &mut AnsweredItem,
1315 (inboxes, commands): (&Inboxes, &mut Commands),
1316 notices: &mut MessageWriter<Notice>,
1317) -> Next {
1318 if let Some(message) = response.message() {
1319 answered.conversation.push(message);
1320 }
1321 let through = answered.conversation.messages().len().checked_sub(1);
1322 let through = through.and_then(|last| answered.conversation.id(last));
1323 if let Some(through) = through.filter(|_| response.usage.context_tokens().is_some()) {
1324 let usage = response.usage;
1325 *answered.last = LastUsage(Some(Anchor { usage, through }));
1326 }
1327 let tool_calls: Vec<ToolCall> = response.tool_calls().cloned().collect();
1328 let failure = turn_failure(
1329 &response.choice,
1330 Some(&response.stop()),
1331 response.finish_reason().as_ref(),
1332 );
1333 if let Some(failure) = failure {
1334 if !tool_calls.is_empty() {
1336 let results = close_pending_with(&tool_calls, &format!("not run: {failure}"));
1337 commit_results(&mut answered.conversation, results);
1338 }
1339 let ending = Ending {
1340 notice: format!("The turn failed: {failure}."),
1341 take_out: true,
1342 };
1343 return Next::End(TurnOutcome::Failed(failure.to_string()), Some(ending));
1344 }
1345 if !tool_calls.is_empty() {
1346 return Next::Run(tool_calls);
1347 }
1348 if let Some(next) = carry_on(inboxes.of(agent)) {
1351 let spec = answered.connection.map(|connection| &*connection.spec);
1352 let conversation = &mut *answered.conversation;
1353 let into = (Some(turn), commands);
1354 deliver(agent, spec, next, conversation, notices, into);
1355 return Next::Call;
1356 }
1357 match answered
1358 .conversation
1359 .messages()
1360 .last()
1361 .map(|message| &**message)
1362 {
1363 Some(answer @ Message::Assistant(_)) => {
1364 Next::End(TurnOutcome::Answered(answer.clone()), None)
1365 }
1366 _ => {
1367 let why = "the model ended the turn without a message".to_owned();
1368 Next::End(TurnOutcome::Failed(why), None)
1369 }
1370 }
1371}
1372
1373fn drop_unanswered(
1377 agent: Entity,
1378 conversation: &mut Conversation,
1379 notices: &mut MessageWriter<Notice>,
1380) {
1381 let unanswered = conversation.messages().last().is_some_and(|message| {
1382 matches!(&**message, Message::User { content }
1383 if !content.iter().any(|item| matches!(item, UserContent::ToolResult(_))))
1384 });
1385 if unanswered {
1386 let last = conversation.messages().len().saturating_sub(1);
1387 conversation.splice(last.., [], Author::rig());
1388 notices.write(Notice::info(
1389 agent,
1390 "Your last message was taken out of the conversation; send it again.",
1391 ));
1392 }
1393}
1394
1395#[derive(QueryData)]
1397#[query_data(mutable)]
1398pub(crate) struct Answered {
1399 conversation: &'static mut Conversation,
1400 connection: Option<&'static Connection>,
1401 last: &'static mut LastUsage,
1402}
1403
1404pub(crate) fn spawn_reply(
1408 commands: &mut Commands,
1409 (agent, turn): (Entity, Entity),
1410 calls: Vec<ToolCall>,
1411 reply: EffectId,
1412 cut_off: &[ToolCall],
1413) {
1414 for call in calls {
1415 let refused = cut_off.iter().find(|cut| cut.id == call.id).map(|cut| {
1416 let name = cut.function.name.as_str();
1417 let why = feedback::output_limit_feedback(name, &cut.function.raw_arguments());
1418 Outcome(Err(ErrorReport::new(
1419 ErrorKind::Tool(ToolErrorKind::InvalidArgs),
1420 why,
1421 )))
1422 });
1423 let mut call = commands.spawn((
1424 ToolRequest { call },
1425 Effect::with_parent_id(Some(reply)),
1426 CallOf(turn),
1427 ForAgent(agent),
1428 ));
1429 if let Some(refused) = refused {
1430 call.insert(refused);
1431 }
1432 }
1433}
1434
1435fn cut_off(response: &CompletionResponse) -> Vec<ToolCall> {
1439 let truncated = response
1440 .finish_reason()
1441 .is_some_and(|finish| finish.truncated_output());
1442 let unfinished = response.unfinished_calls().filter(|_| truncated);
1443 unfinished.cloned().collect()
1444}
1445
1446#[derive(QueryData)]
1448pub(crate) struct ReplyCall {
1449 effect: &'static Effect,
1450 request: &'static ToolRequest,
1451 outcome: Option<&'static Outcome>,
1452 if_stopped: Option<&'static IfStopped>,
1453 successor: Option<&'static SupersededBy>,
1454 footprint: Option<&'static Footprint>,
1455 settled: Has<Settled>,
1456 started: Has<Started>,
1457 waiting: Has<Waiting>,
1458 entity: Entity,
1459}
1460
1461fn reply_calls<'a>(
1464 reply: EffectId,
1465 calls: &Calls,
1466 replies: &'a Query<ReplyCall>,
1467) -> Vec<ReplyCallItem<'a, 'a>> {
1468 let mut asked: Vec<ReplyCallItem<'a, 'a>> = Vec::new();
1469 for call in calls.iter() {
1470 let Ok(mut call) = replies.get(call) else {
1471 continue;
1472 };
1473 let first = |at: &ReplyCallItem| at.request.call.id == call.request.call.id;
1474 if call.effect.parent != Some(reply) || asked.iter().any(first) {
1475 continue;
1476 }
1477 for _ in 0..calls.len() {
1478 match call
1479 .successor
1480 .and_then(|by| replies.get(by.successor()).ok())
1481 {
1482 Some(successor) => call = successor,
1483 None => break,
1484 }
1485 }
1486 asked.push(call);
1487 }
1488 asked
1489}
1490
1491pub(crate) fn wait_in_order(
1498 turns: Query<(&Turn, &Calls)>,
1499 replies: Query<ReplyCall>,
1500 mut commands: Commands,
1501) {
1502 let anything = Footprint::default();
1503 for (state, calls) in &turns {
1504 let Some(reply) = state.reply else {
1505 continue;
1506 };
1507 let mut earlier: Vec<&Footprint> = Vec::new();
1508 for call in reply_calls(reply, calls, &replies) {
1509 let footprint = call.footprint.unwrap_or(&anything);
1510 let waits = earlier
1511 .iter()
1512 .any(|before| footprint.0.conflicts_with(&before.0));
1513 let open = !call.started && call.outcome.is_none() && !call.settled;
1514 if open && waits != call.waiting {
1515 let mut entity = commands.entity(call.entity);
1516 match waits {
1517 true => entity.try_insert(Waiting),
1518 false => entity.try_remove::<Waiting>(),
1519 };
1520 }
1521 if !call.settled {
1522 earlier.push(footprint);
1523 }
1524 }
1525 }
1526}
1527
1528pub(crate) fn commit_replies(
1534 mut turns: Query<(Entity, &TurnOf, &mut Turn, &Calls, Option<&TurnOutcome>)>,
1535 replies: Query<ReplyCall>,
1536 mut agents: Query<&mut Conversation>,
1537 mut commands: Commands,
1538) {
1539 for (turn, &TurnOf(agent), mut state, calls, ending) in &mut turns {
1540 let Some(reply) = state.reply else {
1541 continue;
1542 };
1543 let calls = reply_calls(reply, calls, &replies);
1544 let Ok(mut conversation) = agents.get_mut(agent) else {
1545 continue;
1546 };
1547 if calls.iter().any(|call| !call.settled) {
1548 for call in calls.iter().filter(|call| call.settled) {
1549 if let Some(outcome) = call.outcome {
1550 let result = outcome.tool_result();
1551 mark_settled(&mut conversation, &call.request.call.id, &result);
1552 }
1553 }
1554 continue;
1555 }
1556 state.reply = None;
1557 let unanswered = ToolResult::failed(ToolExecutionError::other("the call had no result"));
1558 if let Some(results) = reply_results(calls, |_| Some(unanswered.clone())) {
1559 commit_results(&mut conversation, results);
1560 }
1561 if ending.and_then(TurnOutcome::stops).is_none() {
1562 spawn_turn_call(&mut commands, turn, agent);
1563 }
1564 }
1565}
1566
1567pub(crate) fn end_turns(
1578 turns: Query<
1579 (
1580 Entity,
1581 &TurnOf,
1582 &TurnOutcome,
1583 Option<&Ending>,
1584 Option<&Calls>,
1585 ),
1586 Without<Holds>,
1587 >,
1588 open: Query<(), (With<Effect>, Without<Settled>)>,
1589 arrived: Query<(), (With<TurnCall>, Added<Settled>)>,
1590 mut agents: Query<(&AgentId, &mut Conversation)>,
1591 inboxes: Inboxes,
1592 mut commands: Commands,
1593 mut notices: MessageWriter<Notice>,
1594) {
1595 for (turn, &TurnOf(agent), outcome, ending, calls) in &turns {
1596 let Ok((id, mut conversation)) = agents.get_mut(agent) else {
1597 continue;
1598 };
1599 if let TurnOutcome::Answered(_) = outcome
1600 && inboxes.coming(agent)
1601 {
1602 continue;
1603 }
1604 let waiting = |call| open.contains(call) || arrived.contains(call);
1605 if calls.is_some_and(|calls| calls.iter().any(waiting)) {
1606 continue;
1607 }
1608 if let Some(ending) = ending {
1609 notices.write(Notice::error(agent, ending.notice.clone()));
1610 if ending.take_out {
1611 drop_unanswered(agent, &mut conversation, &mut notices);
1612 }
1613 }
1614 if !matches!(outcome, TurnOutcome::Answered(_)) {
1615 conversation.halt(Halt::Kept);
1616 }
1617 let end = turn_end(outcome, &conversation);
1618 conversation.end_turn(end);
1619 let how = match outcome {
1620 TurnOutcome::Answered(_) => "answered",
1621 TurnOutcome::Failed(_) => "failed",
1622 TurnOutcome::Stopped(_) => "stopped",
1623 };
1624 info!(agent = %id.0, "turn ended: {how}");
1625 commands
1626 .entity(turn)
1627 .try_remove::<(TurnOf, Ending)>()
1628 .try_insert(Settled);
1629 }
1630}
1631
1632pub(crate) fn carry_on_called(
1636 called: Query<&CallOf, (With<TurnCall>, Added<Admitted>)>,
1637 ending: Query<Option<&Holds>, With<TurnOutcome>>,
1638 mut commands: Commands,
1639) {
1640 for &CallOf(turn) in &called {
1641 if let Ok(holds) = ending.get(turn) {
1642 take_outcome(&mut commands, turn, holds);
1643 }
1644 }
1645}
1646
1647pub(crate) fn take_outcome(commands: &mut Commands, turn: Entity, holds: Option<&Holds>) {
1650 commands.entity(turn).try_remove::<(TurnOutcome, Ending)>();
1651 release(commands, holds);
1652}
1653
1654pub(crate) fn keeps_settled(world: &World, entity: Entity) -> bool {
1657 let reply = world
1658 .get::<CallOf>(entity)
1659 .and_then(|of| world.get::<Turn>(of.0))
1660 .and_then(|turn| turn.reply);
1661 let parent = world.get::<Effect>(entity).and_then(|effect| effect.parent);
1662 reply.is_some() && parent == reply && world.get::<ToolRequest>(entity).is_some()
1663}