Repository navigation
Expand file tree
/
Copy pathagent_loop.rs
More file actions
2979 lines (2782 loc) · 117 KB
/
Copy pathagent_loop.rs
File metadata and controls
2979 lines (2782 loc) · 117 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
//! The core agent loop: prompt → LLM stream → tool execution → repeat.
//!
//! This is the heart of yoagent:
//!
//! - `agent_loop()` starts with new prompt messages
//! - `agent_loop_continue()` resumes from existing context
//!
//! Both return a stream of `AgentEvent`s.
use crate::context::{
self, CompactionStrategy, ContextConfig, ContextTracker, DefaultCompaction, ExecutionLimits,
ExecutionTracker,
};
use crate::provider::{
ModelConfig, ProviderError, StreamConfig, StreamEvent, StreamProvider, ToolDefinition,
};
use crate::types::*;
use std::sync::Arc;
/// Type alias for convert_to_llm callback.
pub type ConvertToLlmFn = Box<dyn Fn(&[AgentMessage]) -> Vec<Message> + Send + Sync>;
/// Type alias for transform_context callback.
pub type TransformContextFn = Box<dyn Fn(Vec<AgentMessage>) -> Vec<AgentMessage> + Send + Sync>;
/// Type alias for steering/follow-up message callbacks.
pub type GetMessagesFn = Box<dyn Fn() -> Vec<AgentMessage> + Send + Sync>;
/// Called before each LLM turn. Return `false` to abort the loop.
pub type BeforeTurnFn = Arc<dyn Fn(&[AgentMessage], usize) -> bool + Send + Sync>;
/// Called after each LLM turn with the current messages and the turn's usage.
pub type AfterTurnFn = Arc<dyn Fn(&[AgentMessage], &Usage) + Send + Sync>;
/// Called when the LLM returns a `StopReason::Error`.
pub type OnErrorFn = Arc<dyn Fn(&str) + Send + Sync>;
use tokio::sync::mpsc;
use tracing::warn;
/// Configuration for the agent loop.
///
/// Build one with [`AgentLoopConfig::new`], then set the fields you need: the
/// struct is `#[non_exhaustive]`, so new fields can be added in a minor release
/// without breaking callers.
///
/// ```
/// use std::sync::Arc;
/// use yoagent::agent_loop::AgentLoopConfig;
/// use yoagent::provider::MockProvider;
///
/// let mut config = AgentLoopConfig::new(Arc::new(MockProvider::text("hi")), "mock");
/// config.max_tokens = Some(1024);
/// ```
#[non_exhaustive]
pub struct AgentLoopConfig {
pub provider: Arc<dyn StreamProvider>,
pub model: String,
pub api_key: String,
pub thinking_level: ThinkingLevel,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
/// Optional model configuration for multi-provider support.
/// When set, passed through to StreamConfig so providers can use
/// base_url, headers, compat flags, etc.
pub model_config: Option<ModelConfig>,
/// Convert AgentMessage[] → Message[] before each LLM call.
/// Default: keep only LLM-compatible messages.
pub convert_to_llm: Option<ConvertToLlmFn>,
/// Transform context before convert_to_llm (for pruning, compaction).
pub transform_context: Option<TransformContextFn>,
/// Get steering messages (user interruptions mid-run).
pub get_steering_messages: Option<GetMessagesFn>,
/// Get follow-up messages (queued work after agent finishes).
pub get_follow_up_messages: Option<GetMessagesFn>,
/// Context window configuration (auto-compaction).
pub context_config: Option<ContextConfig>,
/// Custom compaction strategy. When set, replaces the default
/// `compact_messages()` call. Invoked when `context_config` is `Some`.
pub compaction_strategy: Option<Arc<dyn CompactionStrategy>>,
/// Execution limits (max turns, tokens, duration).
pub execution_limits: Option<ExecutionLimits>,
/// Prompt caching configuration.
pub cache_config: CacheConfig,
/// Where to stash the full text of a truncated tool result, so the model
/// can retrieve what head-tail truncation elided.
///
/// `None` (default) means truncation behaves exactly as before: the middle
/// is gone and only the event stream, which the agent cannot read, still
/// has it. Set it and the marker names a `shared_state` key instead.
///
/// The pointer lives in the transcript, so it is lost when lossy compaction
/// drops that turn — the stash entry outlives it and keeps consuming the
/// backend's cap. Retrieval is best-effort by construction.
pub tool_output_sink: Option<crate::shared_state::SharedState>,
/// Tool execution strategy (sequential, parallel, or batched).
pub tool_execution: ToolExecutionStrategy,
/// Tool middleware chain — approve/deny/modify every tool call before it
/// executes (see [`ToolMiddleware`]). Empty = allow all.
pub tool_middleware: Vec<Arc<dyn ToolMiddleware>>,
/// Structured-output constraint, passed through to the provider (see
/// [`OutputSchema`](crate::provider::OutputSchema)). Usually set via
/// [`Agent::prompt_structured`](crate::Agent::prompt_structured).
pub output_schema: Option<crate::provider::OutputSchema>,
/// Retry configuration for transient provider errors.
pub retry_config: crate::retry::RetryConfig,
/// Called before each LLM turn. Return `false` to abort the loop.
pub before_turn: Option<BeforeTurnFn>,
/// Called after each LLM turn with the current messages and the turn's usage.
pub after_turn: Option<AfterTurnFn>,
/// Called when the LLM returns a `StopReason::Error`.
pub on_error: Option<OnErrorFn>,
/// Input filters applied to user messages before the LLM call.
/// Filters run in order; first `Reject` wins and discards any accumulated
/// warnings. `Warn` messages accumulate and are appended to the user message.
/// An [`AsyncFilter`] in the list is awaited (see [`InputFilter::as_async`]).
pub input_filters: Vec<Arc<dyn InputFilter>>,
/// Optional delay between turns. Useful for rate-limit-sensitive scenarios
/// (e.g., OAuth tokens with low request-per-minute caps). Skipped on the
/// first turn so the agent starts immediately.
pub turn_delay: Option<std::time::Duration>,
/// Extensions for this run (see [`crate::extension`]), in dispatch order
/// after `tree_extensions`.
pub extensions: Vec<Arc<dyn crate::extension::Extension>>,
/// Extensions for this run and every run it delegates to, at any depth
/// (host policy). They run before `extensions`.
pub tree_extensions: Vec<Arc<dyn crate::extension::Extension>>,
/// How many times per run an extension's `on_stop` may continue it.
pub max_stop_continues: usize,
/// The host's label for the run, passed to extensions as
/// [`RunContext::label`](crate::extension::RunContext::label).
pub run_label: Option<String>,
/// Tree extensions a parent run passed down. They run first, and their
/// tools are not offered.
pub(crate) inherited_extensions: Vec<Arc<dyn crate::extension::Extension>>,
/// 0 for a top-level run; a delegated run's depth.
pub(crate) depth: usize,
/// The tool call that started this run, for a delegated run.
pub(crate) delegated_by: Option<String>,
/// The calling run's extension run id, for a delegated run.
pub(crate) parent_run_id: Option<String>,
}
impl AgentLoopConfig {
/// A configuration for `model` on `provider`, with every other field at
/// its default: no API key, thinking off, no limits beyond the provider's,
/// no context management, no hooks, parallel tool execution and the
/// default retry policy.
pub fn new(provider: Arc<dyn StreamProvider>, model: impl Into<String>) -> Self {
Self {
provider,
model: model.into(),
api_key: String::new(),
thinking_level: ThinkingLevel::Off,
max_tokens: None,
temperature: None,
model_config: None,
convert_to_llm: None,
transform_context: None,
get_steering_messages: None,
get_follow_up_messages: None,
context_config: None,
compaction_strategy: None,
execution_limits: None,
cache_config: CacheConfig::default(),
tool_output_sink: None,
tool_execution: ToolExecutionStrategy::default(),
tool_middleware: Vec::new(),
output_schema: None,
retry_config: crate::retry::RetryConfig::default(),
before_turn: None,
after_turn: None,
on_error: None,
input_filters: Vec::new(),
turn_delay: None,
extensions: Vec::new(),
tree_extensions: Vec::new(),
max_stop_continues: crate::extension::DEFAULT_MAX_STOP_CONTINUES,
run_label: None,
inherited_extensions: Vec::new(),
depth: 0,
delegated_by: None,
parent_run_id: None,
}
}
/// Whether any extension applies to this run.
fn has_extensions(&self) -> bool {
!(self.extensions.is_empty()
&& self.tree_extensions.is_empty()
&& self.inherited_extensions.is_empty())
}
/// Make this the configuration of a run delegated by the tool call
/// `ctx` belongs to: the calling run's tree extensions apply here (ahead
/// of this run's own, their tools not offered), at the delegation's
/// depth, under the calling run's label. What a custom delegation tool
/// calls before running its child; `SubAgentTool` does it itself.
pub fn delegated_from(&mut self, ctx: &ToolContext) -> &mut Self {
self.inherited_extensions = ctx.tree_extensions().to_vec();
self.depth = ctx.delegation_depth();
self.delegated_by = ctx.delegation.call_id.clone();
self.parent_run_id = ctx.delegation.parent_run_id.clone();
if self.run_label.is_none() {
self.run_label = ctx.run_label().map(String::from);
}
self
}
/// The tree extensions a delegated run inherits: this run's inherited
/// ones, then its own.
pub(crate) fn tree_for_children(&self) -> Vec<Arc<dyn crate::extension::Extension>> {
self.inherited_extensions
.iter()
.chain(&self.tree_extensions)
.cloned()
.collect()
}
}
/// Default convert_to_llm: keep only user/assistant/toolResult messages.
fn default_convert_to_llm(messages: &[AgentMessage]) -> Vec<Message> {
messages
.iter()
.filter_map(|m| m.as_llm().cloned())
.collect()
}
/// Start an agent loop with new prompt messages.
/// Prefix of the message the loop appends when it stops a run itself.
///
/// A run stopped by a limit or by loop detection ends with a `Message::User`
/// carrying this prefix. The last *assistant* message still reports whatever
/// stop reason it had — `ToolUse` for a loop abort — so a consumer inspecting
/// only assistant messages cannot tell a stopped run from a finished one.
/// `SubAgentTool` uses this to avoid reporting an aborted delegation as a
/// bland success.
pub const AGENT_STOPPED_PREFIX: &str = "[Agent stopped:";
/// The stop marker for a run cancelled between provider calls — while tools
/// ran, or between turns. (A cancel during a provider call ends that turn's
/// message as `StopReason::Aborted` instead.) Unlike a limit, which leaves
/// real partial work, a cancelled run did not finish: `SubAgentTool` reports
/// it as a failure.
pub const CANCELLED_MARKER: &str = "[Agent stopped: cancelled]";
/// End a turn that stops before its model request, so its `TurnStart` is
/// paired. It made no assistant message and ran no tools: `TurnEnd` carries
/// the last message of the history (the stop marker, when there is one).
/// With an empty history (a run given no prompts at all) there is no message
/// to carry, and the turn is left open.
fn close_turn(tx: &mpsc::UnboundedSender<AgentEvent>, context: &AgentContext) {
if let Some(last) = context.messages.last() {
tx.send(AgentEvent::TurnEnd {
message: last.clone(),
tool_results: Vec::new(),
})
.ok();
}
}
/// Append a stop marker as a user message: emitted, kept in context and in
/// the run's new messages.
fn push_stop_marker(
text: String,
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
) {
let marker = AgentMessage::Llm(Message::User {
content: vec![Content::Text { text }],
timestamp: now_ms(),
});
tx.send(AgentEvent::MessageStart {
message: marker.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: marker.clone(),
})
.ok();
context.messages.push(marker.clone());
new_messages.push(marker);
}
/// Mark a run cancelled between provider calls — only when it did something,
/// so a run cancelled before it started leaves no trace.
fn mark_cancelled(
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
) {
if new_messages
.iter()
.any(|m| matches!(m, AgentMessage::Llm(Message::Assistant { .. })))
{
push_stop_marker(CANCELLED_MARKER.to_string(), tx, context, new_messages);
}
}
/// The stop marker for a run halted by loop detection specifically.
///
/// Distinct from the limit stops because the two mean opposite things to a
/// caller. Hitting `max_turns` is a bound: the work was cut short but what it
/// produced is real, and `SubAgentTool` returns it. A loop abort is a failure:
/// the model was emitting the same call forever and there is nothing to keep.
pub const LOOP_ABORT_PREFIX: &str = "[Agent stopped: repeated tool call —";
/// Prefix of the user-role nudge the loop injects when loop detection steers
/// a model that keeps repeating one call.
pub(crate) const LOOP_NUDGE_PREFIX: &str = "[You have called ";
/// Per-run state the loop shares with the hooks that run inside its task.
#[derive(Default)]
struct LoopScope {
/// Decision-model spend, recorded by the decision integrations and
/// folded into the run's `SessionStats`.
decision: DecisionStats,
/// Compaction-model spend, recorded by `LlmCompaction::compact` (which
/// the loop calls synchronously inside its task) and folded into the
/// run's `SessionStats::compaction`.
compaction: CompactionSpend,
/// The user messages this run was given — its prompts, steering and
/// follow-ups — kept apart from the context so compaction cannot remove
/// them. Exposed as `ToolCallRequest::run_prompts` /
/// `TurnContext::run_prompts`.
prompts: Vec<Message>,
/// The same decision and compaction spend, piece by piece with each
/// piece's own price, for the run's extensions (`Budget`).
spend: Option<Arc<crate::extension::RunSpend>>,
}
tokio::task_local! {
/// The running loop's [`LoopScope`]. Hooks run inside the loop's task; a
/// nested sub-agent loop opens its own scope.
static LOOP_SCOPE: std::cell::RefCell<LoopScope>;
}
/// Record decision-model spend into the enclosing loop's stats. A no-op
/// outside a loop.
#[cfg_attr(not(feature = "decision"), allow(dead_code))]
pub(crate) fn record_decision(f: impl FnOnce(&mut DecisionStats)) {
let _ = LOOP_SCOPE.try_with(|cell| {
// Recorded on its own first, so the piece's own price is known.
let mut piece = DecisionStats::default();
f(&mut piece);
let mut scope = cell.borrow_mut();
if let Some(spend) = &scope.spend {
spend.add(&piece.usage, piece.cost_usd);
}
scope.decision.merge(&piece);
});
}
/// Record compaction-model spend into the enclosing loop's stats. `f` runs
/// only inside a loop, so a caller can move spend out of its own ledger in
/// `f` without losing it outside one; returns whether it ran.
pub(crate) fn record_compaction(f: impl FnOnce(&mut CompactionSpend)) -> bool {
LOOP_SCOPE
.try_with(|cell| {
let mut piece = CompactionSpend::default();
f(&mut piece);
let mut scope = cell.borrow_mut();
if let Some(spend) = &scope.spend {
spend.add(&piece.usage, piece.cost_usd);
}
scope.compaction.merge(&piece);
})
.is_ok()
}
/// What hooks inside a run recorded into its [`LoopScope`].
#[derive(Default)]
struct ScopeSpend {
decision: DecisionStats,
compaction: CompactionSpend,
}
impl ScopeSpend {
/// Fold into the run's stats.
fn fold_into(&self, stats: &mut SessionStats) {
stats.decision.merge(&self.decision);
stats.compaction.merge(&self.compaction);
}
}
/// The user messages the enclosing loop's run was given so far (empty
/// outside a loop).
pub(crate) fn run_prompts() -> Vec<Message> {
LOOP_SCOPE
.try_with(|cell| cell.borrow().prompts.clone())
.unwrap_or_default()
}
/// The user-role LLM messages among `messages`.
fn user_messages(messages: &[AgentMessage]) -> Vec<Message> {
messages
.iter()
.filter_map(|m| match m {
AgentMessage::Llm(msg @ Message::User { .. }) => Some(msg.clone()),
_ => None,
})
.collect()
}
/// Note messages handed to the run (steering, follow-ups) as run prompts.
fn note_run_prompts(messages: &[AgentMessage]) {
let users = user_messages(messages);
if !users.is_empty() {
let _ = LOOP_SCOPE.try_with(|cell| cell.borrow_mut().prompts.extend(users));
}
}
/// Run `fut` in a fresh [`LoopScope`] seeded with the run's prompts, and
/// return its output plus the spend recorded.
async fn with_loop_scope<T>(
prompts: Vec<Message>,
spend: Arc<crate::extension::RunSpend>,
fut: impl std::future::Future<Output = T>,
) -> (T, ScopeSpend) {
let scope = LoopScope {
prompts,
spend: Some(spend),
..Default::default()
};
LOOP_SCOPE
.scope(std::cell::RefCell::new(scope), async {
let out = fut.await;
let recorded = LOOP_SCOPE.with(|cell| {
let mut scope = cell.borrow_mut();
ScopeSpend {
decision: std::mem::take(&mut scope.decision),
compaction: std::mem::take(&mut scope.compaction),
}
});
(out, recorded)
})
.await
}
/// The error tool result given to each tool call in a response that ended as
/// [`StopReason::Refusal`] — the model declined, or a content filter stopped
/// the response. The call is never executed.
const REFUSAL_TOOL_RESULT_TEXT: &str = "Tool call not run: the response was stopped as a refusal \
(declined by the model or stopped by the content filter).";
/// The error tool result given to a tool call whose run was already
/// cancelled when the call was reached. The call is never executed.
const CANCELLED_TOOL_RESULT_TEXT: &str = "Tool call not run: the run was cancelled.";
pub async fn agent_loop(
prompts: Vec<AgentMessage>,
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> Vec<AgentMessage> {
agent_loop_with_stats(prompts, context, config, tx, cancel)
.await
.0
}
/// [`agent_loop`], also returning the [`SessionStats`] it sent on
/// `AgentEnd` — for callers inside the crate that must not depend on someone
/// draining the event channel (`Agent`'s sub-agent bucket, `SubAgentTool`'s
/// spend report).
pub(crate) async fn agent_loop_with_stats(
prompts: Vec<AgentMessage>,
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> (Vec<AgentMessage>, SessionStats) {
let prompt_messages = user_messages(&prompts);
let (exts, start_failure) = start_extensions(config, &prompt_messages, &cancel).await;
let observer = EventObserver::new(&exts, tx);
let tx = observer.sender();
tx.send(AgentEvent::AgentStart).ok();
// One scope for the whole run, input filters included, so decision
// spend in an async filter is counted — also when it rejects.
let (outcome, recorded) = with_loop_scope(Vec::new(), exts.spend(), async {
// What the filters see, without the warnings they append.
let text = if exts.is_empty() {
String::new()
} else {
prompt_text(&prompts)
};
let prompts = apply_input_filters(prompts, config).await?;
if !exts.is_empty() {
exts.on_input(&crate::extension::InputContext::new(&text, &prompts))
.await?;
}
note_run_prompts(&prompts);
let mut new_messages: Vec<AgentMessage> = prompts.clone();
// Add prompts to context
for prompt in &prompts {
context.messages.push(prompt.clone());
}
tx.send(AgentEvent::TurnStart).ok();
// Emit events for each prompt message
for prompt in &prompts {
tx.send(AgentEvent::MessageStart {
message: prompt.clone(),
})
.ok();
tx.send(AgentEvent::MessageEnd {
message: prompt.clone(),
})
.ok();
}
let stats = run_with_extensions(
context,
&mut new_messages,
config,
&tx,
&cancel,
&exts,
start_failure,
)
.await;
Ok::<_, String>((new_messages, stats))
})
.await;
let (new_messages, mut stats) = match outcome {
Ok(done) => done,
Err(reason) => {
tx.send(AgentEvent::InputRejected {
reason: reason.clone(),
})
.ok();
exts.finish(&crate::extension::RunOutcome::new(
crate::extension::RunEnd::Rejected {
reason: reason.clone(),
},
None,
))
.await;
let mut stats = SessionStats::default();
recorded.fold_into(&mut stats);
tx.send(AgentEvent::AgentEnd {
messages: vec![],
stats: stats.clone(),
})
.ok();
drop(tx);
observer.finish().await;
return (vec![], stats);
}
};
recorded.fold_into(&mut stats);
exts.finish(&run_outcome(&new_messages, &cancel)).await;
tx.send(AgentEvent::AgentEnd {
messages: new_messages.clone(),
stats: stats.clone(),
})
.ok();
drop(tx);
observer.finish().await;
(new_messages, stats)
}
/// Start the run's extensions. A required extension (or one that filters
/// tool output) that cannot start makes the run fail once it has begun; the
/// extensions that did start are returned with the failure, so they still
/// get `finish`.
async fn start_extensions(
config: &AgentLoopConfig,
prompts: &[Message],
cancel: &tokio_util::sync::CancellationToken,
) -> (
Arc<crate::extension::ActiveExtensions>,
Option<crate::extension::Failure>,
) {
use crate::extension::{ActiveExtensions, RunInfo};
if !config.has_extensions() {
return (Arc::new(ActiveExtensions::none()), None);
}
let info = RunInfo {
label: config.run_label.as_deref(),
prompts,
depth: config.depth,
delegated_by: config.delegated_by.as_deref(),
parent_run_id: config.parent_run_id.as_deref(),
cancel,
};
let own: Vec<_> = config
.tree_extensions
.iter()
.chain(&config.extensions)
.cloned()
.collect();
let (exts, failure) = ActiveExtensions::start(&config.inherited_extensions, &own, info).await;
(Arc::new(exts), failure)
}
/// Run the loop with the run's extensions: offer their tools for this run
/// only, or fail the run when a required extension could not start.
#[allow(clippy::too_many_arguments)]
async fn run_with_extensions(
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
cancel: &tokio_util::sync::CancellationToken,
exts: &crate::extension::ActiveExtensions,
start_failure: Option<crate::extension::Failure>,
) -> SessionStats {
if let Some(failure) = start_failure {
fail_run(&failure, exts, config, tx, context, new_messages, true);
return SessionStats::default();
}
let base_tools = context.tools.len();
if !exts.is_empty() {
let tools = exts.tools().await;
if !tools.is_empty() {
crate::tool_source::merge(&mut context.tools, tools);
}
}
let stats = {
use tracing::Instrument;
run_loop(context, new_messages, config, tx, cancel, exts)
.instrument(tracing::info_span!("agent_loop", model = %config.model))
.await
};
// Extension tools belong to this run only.
context.tools.truncate(base_tools);
// However the loop ended (a limit, a cancel, a stop, a provider error), a
// required extension's failure it did not act on still fails the run.
if let Some(failure) = exts.settle().await {
fail_run(&failure, exts, config, tx, context, new_messages, false);
}
stats
}
/// End the run because a required extension failed: an assistant error
/// message naming it, announced and appended, `on_error`, inside a turn. With
/// `in_turn`, the current turn is closed with it; otherwise it gets a turn of
/// its own, so every message stays between a `TurnStart` and a `TurnEnd`.
fn fail_run(
failure: &crate::extension::Failure,
exts: &crate::extension::ActiveExtensions,
config: &AgentLoopConfig,
tx: &mpsc::UnboundedSender<AgentEvent>,
context: &mut AgentContext,
new_messages: &mut Vec<AgentMessage>,
in_turn: bool,
) {
// The run fails once: a failure recorded while this one is reported
// (an `on_event` that fails on these very events) is only logged.
exts.latch_failed();
let text = failure.message();
tracing::error!("{text}");
if !in_turn {
tx.send(AgentEvent::TurnStart).ok();
}
let message = error_message(&config.model, text.clone());
announce(tx, &message);
let am: AgentMessage = message.into();
context.messages.push(am.clone());
new_messages.push(am.clone());
if let Some(ref on_error) = config.on_error {
on_error(&text);
}
tx.send(AgentEvent::TurnEnd {
message: am,
tool_results: Vec::new(),
})
.ok();
}
/// How a run ended, from its new messages.
fn run_outcome(
new_messages: &[AgentMessage],
cancel: &tokio_util::sync::CancellationToken,
) -> crate::extension::RunOutcome {
use crate::extension::{failed_extension, RunEnd, RunOutcome};
let stop_reason = new_messages.iter().rev().find_map(|m| match m {
AgentMessage::Llm(Message::Assistant { stop_reason, .. }) => Some(stop_reason.clone()),
_ => None,
});
let end = match new_messages.last() {
// A stop marker ends the run: a limit, loop detection, a cancel
// between turns, `on_before_turn`, an extension's `Stop`.
Some(AgentMessage::Llm(Message::User { content, .. })) => match content.first() {
Some(Content::Text { text }) if text == CANCELLED_MARKER => RunEnd::Cancelled,
Some(Content::Text { text }) if text.starts_with(AGENT_STOPPED_PREFIX) => {
RunEnd::Stopped {
reason: text.clone(),
}
}
_ if cancel.is_cancelled() => RunEnd::Cancelled,
_ => RunEnd::Completed,
},
Some(AgentMessage::Llm(Message::Assistant {
stop_reason: StopReason::Aborted,
..
})) => RunEnd::Cancelled,
// The model finished (or stopped on its own): a cancel after that
// changes nothing.
Some(AgentMessage::Llm(Message::Assistant {
stop_reason: StopReason::Stop | StopReason::Length | StopReason::Refusal,
..
})) => RunEnd::Completed,
Some(AgentMessage::Llm(Message::Assistant {
stop_reason: StopReason::Error,
error_message,
..
})) => {
let error = error_message.clone().unwrap_or_default();
RunEnd::Failed {
extension: failed_extension(&error),
error,
}
}
_ if cancel.is_cancelled() => RunEnd::Cancelled,
_ => RunEnd::Completed,
};
RunOutcome::new(end, stop_reason)
}
/// Every event of a run with extensions passes through their `on_event`
/// before it reaches the run's consumer. Without extensions, the consumer's
/// sender is used as is. The loop waits on a flush before each failure check
/// and before `finish`, so what `on_event` recorded is current there.
/// Failures on the run's last events (`AgentEnd` itself) come too late to
/// change its outcome and are only logged. Events a stray sender clone sends
/// after the run ended are dropped.
struct EventObserver {
tx: mpsc::UnboundedSender<AgentEvent>,
exts: Arc<crate::extension::ActiveExtensions>,
task: Option<(tokio::sync::oneshot::Sender<()>, crate::rt::JoinHandle<()>)>,
}
impl EventObserver {
fn new(
exts: &Arc<crate::extension::ActiveExtensions>,
out: mpsc::UnboundedSender<AgentEvent>,
) -> Self {
if exts.is_empty() {
return Self {
tx: out,
exts: exts.clone(),
task: None,
};
}
let (tx, mut rx) = mpsc::unbounded_channel::<AgentEvent>();
let (flush_tx, mut flush_rx) = mpsc::unbounded_channel::<crate::extension::FlushRequest>();
exts.connect_observer(flush_tx);
let (done_tx, mut done_rx) = tokio::sync::oneshot::channel::<()>();
let observed = exts.clone();
// What the hooks log goes where the run's own logs go.
use tracing::instrument::{Instrument, WithSubscriber};
let task = crate::rt::spawn(
async move {
loop {
tokio::select! {
biased;
event = rx.recv() => match event {
Some(event) => {
observed.on_event(&event).await;
out.send(event).ok();
}
None => break,
},
Some(ack) = flush_rx.recv() => {
// Every event sent before the request is queued.
while let Ok(event) = rx.try_recv() {
observed.on_event(&event).await;
out.send(event).ok();
}
let _ = ack.send(());
}
_ = &mut done_rx => {
// The run is over: deliver what is queued and stop,
// even if a stray sender clone is still alive.
while let Ok(event) = rx.try_recv() {
observed.on_event(&event).await;
out.send(event).ok();
}
break;
}
}
}
}
.in_current_span()
.with_current_subscriber(),
);
Self {
tx,
exts: exts.clone(),
task: Some((done_tx, task)),
}
}
fn sender(&self) -> mpsc::UnboundedSender<AgentEvent> {
self.tx.clone()
}
/// Deliver every event sent so far, then stop observing.
async fn finish(self) {
let Self { tx, exts, task } = self;
drop(tx);
if let Some((done, task)) = task {
let _ = done.send(());
if let Err(e) = task.await {
tracing::error!("extension event observer failed: {e}");
}
}
if let Some(failure) = exts.take_failure() {
tracing::error!(
"{} (on the run's last events: too late to change its outcome)",
failure.message()
);
}
}
}
/// Every user text block of `prompts`, joined by newlines.
fn prompt_text(prompts: &[AgentMessage]) -> String {
prompts
.iter()
.filter_map(|m| {
if let AgentMessage::Llm(Message::User { content, .. }) = m {
Some(
content
.iter()
.filter_map(|c| {
if let Content::Text { text } = c {
Some(text.as_str())
} else {
None
}
})
.collect::<Vec<_>>()
.join("\n"),
)
} else {
None
}
})
.collect::<Vec<_>>()
.join("\n")
}
/// Run the input filters over the prompts: `Err(reason)` on the first
/// reject, otherwise the prompts with any warnings appended to the last
/// user message.
async fn apply_input_filters(
prompts: Vec<AgentMessage>,
config: &AgentLoopConfig,
) -> Result<Vec<AgentMessage>, String> {
let prompts = if !config.input_filters.is_empty() {
let user_text = prompt_text(&prompts);
let mut warnings: Vec<String> = Vec::new();
for filter in &config.input_filters {
let verdict = match filter.as_async() {
// A panicking async filter must not take the loop task (and
// with it the agent's tools and history) down: contain it and
// fail closed, as for middleware.
Some(async_filter) => {
use futures::FutureExt;
std::panic::AssertUnwindSafe(async_filter.filter(&user_text))
.catch_unwind()
.await
.unwrap_or_else(|_| {
warn!("async input filter panicked; rejecting the input");
FilterResult::Reject("input filter panicked".into())
})
}
None => filter.filter(&user_text),
};
match verdict {
FilterResult::Pass => {}
FilterResult::Warn(w) => warnings.push(w),
FilterResult::Reject(reason) => return Err(reason),
}
}
// Append warnings to the last user message's content (avoids consecutive user messages)
if !warnings.is_empty() {
let warning_text = warnings
.iter()
.map(|w| format!("[Warning: {}]", w))
.collect::<Vec<_>>()
.join("\n");
let mut modified = prompts;
// Find and extend the last user message
for msg in modified.iter_mut().rev() {
if let AgentMessage::Llm(Message::User { content, .. }) = msg {
content.push(Content::Text { text: warning_text });
break;
}
}
modified
} else {
prompts
}
} else {
prompts
};
Ok(prompts)
}
/// Continue an agent loop from existing context (for retries).
pub async fn agent_loop_continue(
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> Vec<AgentMessage> {
agent_loop_continue_with_stats(context, config, tx, cancel)
.await
.0
}
/// [`agent_loop_continue`], also returning its [`SessionStats`].
pub(crate) async fn agent_loop_continue_with_stats(
context: &mut AgentContext,
config: &AgentLoopConfig,
tx: mpsc::UnboundedSender<AgentEvent>,
cancel: tokio_util::sync::CancellationToken,
) -> (Vec<AgentMessage>, SessionStats) {
assert!(
!context.messages.is_empty(),
"Cannot continue: no messages in context"
);
if let Some(last) = context.messages.last() {
assert!(
last.role() != "assistant",
"Cannot continue from assistant message"
);
}
let (exts, start_failure) = start_extensions(config, &[], &cancel).await;
let observer = EventObserver::new(&exts, tx);
let tx = observer.sender();
let mut new_messages: Vec<AgentMessage> = Vec::new();
tx.send(AgentEvent::AgentStart).ok();
tx.send(AgentEvent::TurnStart).ok();
let stats = {
let (mut stats, recorded) = with_loop_scope(
Vec::new(),
exts.spend(),
run_with_extensions(
context,
&mut new_messages,
config,
&tx,
&cancel,
&exts,
start_failure,
),
)
.await;
recorded.fold_into(&mut stats);
stats
};
exts.finish(&run_outcome(&new_messages, &cancel)).await;
tx.send(AgentEvent::AgentEnd {
messages: new_messages.clone(),
stats: stats.clone(),
})
.ok();
drop(tx);
observer.finish().await;
(new_messages, stats)
}