1use std::sync::atomic::{AtomicU64, Ordering};
67use std::sync::Arc;
68
69use tokio::sync::mpsc;
70
71use crate::extensions::commands::CommandOutputEvent;
72use crate::extensions::runtime::InvokeCommandEvent;
73use crate::extensions::tasks::TaskEvent;
74
75pub const INVOKE_EVENT_QUEUE_CAPACITY: usize = 8;
80
81pub const INVOKE_RETAINED_BYTE_BUDGET: usize = 256 * 1024;
85
86pub const INVOKE_RETAINED_EVENT_BUDGET: usize = 1024;
90
91pub const INVOKE_CONTROL_RESERVE_BYTES: usize = 16 * 1024;
97
98#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub struct InvokeOutputBudget {
101 pub max_retained_payload_bytes: usize,
104 pub max_retained_events: usize,
107 pub control_reserve_bytes: usize,
109}
110
111impl Default for InvokeOutputBudget {
112 fn default() -> Self {
113 Self {
114 max_retained_payload_bytes: INVOKE_RETAINED_BYTE_BUDGET,
115 max_retained_events: INVOKE_RETAINED_EVENT_BUDGET,
116 control_reserve_bytes: INVOKE_CONTROL_RESERVE_BYTES,
117 }
118 }
119}
120
121impl InvokeOutputBudget {
122 pub fn unlimited() -> Self {
127 Self {
128 max_retained_payload_bytes: usize::MAX,
129 max_retained_events: usize::MAX,
130 control_reserve_bytes: 0,
131 }
132 }
133}
134
135#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
137pub struct InvokeOutputCounters {
138 pub produced_events: u64,
141 pub produced_payload_bytes: u64,
142 pub consumed_events: u64,
144 pub consumed_payload_bytes: u64,
145 pub retained_events: u64,
148 pub retained_payload_bytes: u64,
149 pub truncated_events: u64,
152 pub truncated_payload_bytes: u64,
153 pub dropped_events: u64,
155 pub dropped_payload_bytes: u64,
156 pub dropped_error_events: u64,
159 pub saw_done: bool,
162}
163
164#[derive(Debug)]
167pub struct InvokeOutputReport {
168 pub events: Vec<InvokeCommandEvent>,
170 pub counters: InvokeOutputCounters,
171}
172
173impl InvokeOutputReport {
174 pub fn is_limited(&self) -> bool {
176 self.counters.dropped_events > 0 || self.counters.truncated_events > 0
177 }
178
179 pub fn limit_notice(&self) -> Option<InvokeOutputNotice> {
184 if !self.is_limited() {
185 return None;
186 }
187 let c = &self.counters;
188 let mut message = format!(
189 "command output limited by invocation budget: produced {} bytes in {} events; \
190 retained {} bytes in {} events; truncated {} bytes across {} events; \
191 dropped {} bytes across {} events",
192 c.produced_payload_bytes,
193 c.produced_events,
194 c.retained_payload_bytes,
195 c.retained_events,
196 c.truncated_payload_bytes,
197 c.truncated_events,
198 c.dropped_payload_bytes,
199 c.dropped_events,
200 );
201 if c.dropped_error_events > 0 {
202 message.push_str(&format!(
203 " (including {} error events)",
204 c.dropped_error_events
205 ));
206 }
207 Some(InvokeOutputNotice {
208 severity_error: c.dropped_error_events > 0,
209 message,
210 })
211 }
212}
213
214#[derive(Debug, Clone, PartialEq, Eq)]
216pub struct InvokeOutputNotice {
217 pub message: String,
218 pub severity_error: bool,
220}
221
222pub fn event_payload_bytes(event: &InvokeCommandEvent) -> usize {
225 match event {
226 InvokeCommandEvent::Output(CommandOutputEvent::Text { content })
227 | InvokeCommandEvent::Output(CommandOutputEvent::System { content })
228 | InvokeCommandEvent::Output(CommandOutputEvent::Error { content }) => content.len(),
229 InvokeCommandEvent::Output(CommandOutputEvent::Table { headers, rows }) => {
230 headers.iter().map(String::len).sum::<usize>()
231 + rows
232 .iter()
233 .map(|r| r.iter().map(String::len).sum::<usize>())
234 .sum::<usize>()
235 }
236 InvokeCommandEvent::Output(CommandOutputEvent::Done) => 0,
237 InvokeCommandEvent::Task(TaskEvent::Start { id, label, .. }) => id.len() + label.len(),
238 InvokeCommandEvent::Task(TaskEvent::Update { id, message, .. }) => {
239 id.len() + message.as_deref().map(str::len).unwrap_or(0)
240 }
241 InvokeCommandEvent::Task(TaskEvent::Log { id, line }) => id.len() + line.len(),
242 InvokeCommandEvent::Task(TaskEvent::Done { id, error }) => {
243 id.len() + error.as_deref().map(str::len).unwrap_or(0)
244 }
245 }
246}
247
248#[derive(Debug, Default)]
250struct InvokeMeter {
251 produced_events: AtomicU64,
252 produced_payload_bytes: AtomicU64,
253}
254
255#[derive(Debug, Clone, Copy, PartialEq, Eq)]
259pub struct InvokeSinkClosed;
260
261#[derive(Debug)]
267pub struct InvokeEventSink {
268 tx: mpsc::Sender<InvokeCommandEvent>,
269 meter: Arc<InvokeMeter>,
270}
271
272impl InvokeEventSink {
273 pub async fn send(&self, event: InvokeCommandEvent) -> Result<(), InvokeSinkClosed> {
276 self.meter.produced_events.fetch_add(1, Ordering::SeqCst);
277 self.meter
278 .produced_payload_bytes
279 .fetch_add(event_payload_bytes(&event) as u64, Ordering::SeqCst);
280 self.tx.send(event).await.map_err(|_| InvokeSinkClosed)
281 }
282
283 pub fn is_closed(&self) -> bool {
285 self.tx.is_closed()
286 }
287}
288
289#[derive(Debug)]
291pub struct InvokeOutputCollector {
292 rx: mpsc::Receiver<InvokeCommandEvent>,
293 meter: Arc<InvokeMeter>,
294 budget: InvokeOutputBudget,
295}
296
297pub fn invoke_event_channel(
299 budget: InvokeOutputBudget,
300) -> (InvokeEventSink, InvokeOutputCollector) {
301 let (tx, rx) = mpsc::channel(INVOKE_EVENT_QUEUE_CAPACITY);
302 let meter = Arc::new(InvokeMeter::default());
303 (
304 InvokeEventSink {
305 tx,
306 meter: Arc::clone(&meter),
307 },
308 InvokeOutputCollector { rx, meter, budget },
309 )
310}
311
312impl InvokeOutputCollector {
313 pub async fn collect(mut self) -> InvokeOutputReport {
318 let mut state = CollectorState::new(self.budget);
319 while let Some(event) = self.rx.recv().await {
320 state.ingest(event);
321 }
322 state.finish(&self.meter)
323 }
324}
325
326struct CollectorState {
329 budget: InvokeOutputBudget,
330 events: Vec<InvokeCommandEvent>,
331 counters: InvokeOutputCounters,
332 ordinary_bytes: usize,
334 control_bytes: usize,
336 started_task_ids: std::collections::HashSet<String>,
338 done_task_ids: std::collections::HashSet<String>,
340}
341
342impl CollectorState {
343 fn new(budget: InvokeOutputBudget) -> Self {
344 Self {
345 budget,
346 events: Vec::new(),
347 counters: InvokeOutputCounters::default(),
348 ordinary_bytes: 0,
349 control_bytes: 0,
350 started_task_ids: std::collections::HashSet::new(),
351 done_task_ids: std::collections::HashSet::new(),
352 }
353 }
354
355 fn ingest(&mut self, event: InvokeCommandEvent) {
356 let bytes = event_payload_bytes(&event);
357 self.counters.consumed_events += 1;
358 self.counters.consumed_payload_bytes += bytes as u64;
359
360 match event {
361 InvokeCommandEvent::Output(CommandOutputEvent::Done) => {
362 if self.counters.saw_done {
363 self.drop_whole(0, false);
365 } else {
366 self.counters.saw_done = true;
367 self.retain_control(InvokeCommandEvent::Output(CommandOutputEvent::Done), 0);
369 }
370 }
371 InvokeCommandEvent::Task(TaskEvent::Done { id, error }) => {
372 let within_ordinary = self.fits_ordinary(bytes);
373 let already_done = self.done_task_ids.contains(&id);
374 if within_ordinary && !already_done {
375 self.done_task_ids.insert(id.clone());
376 self.retain_ordinary(
377 InvokeCommandEvent::Task(TaskEvent::Done { id, error }),
378 bytes,
379 );
380 } else if !already_done
381 && self.started_task_ids.contains(&id)
382 && self.control_bytes + bytes <= self.budget.control_reserve_bytes
383 {
384 self.done_task_ids.insert(id.clone());
387 self.retain_control(
388 InvokeCommandEvent::Task(TaskEvent::Done { id, error }),
389 bytes,
390 );
391 } else {
392 self.drop_whole(bytes, false);
393 }
394 }
395 event => self.ingest_ordinary(event, bytes),
396 }
397 }
398
399 fn fits_ordinary(&self, bytes: usize) -> bool {
400 (self.counters.retained_events as usize) < self.budget.max_retained_events
401 && self
402 .ordinary_bytes
403 .checked_add(bytes)
404 .is_some_and(|total| total <= self.budget.max_retained_payload_bytes)
405 }
406
407 fn ingest_ordinary(&mut self, event: InvokeCommandEvent, bytes: usize) {
408 let is_error = is_error_event(&event);
409 if (self.counters.retained_events as usize) >= self.budget.max_retained_events {
410 self.drop_whole(bytes, is_error);
411 return;
412 }
413 let remaining = self
414 .budget
415 .max_retained_payload_bytes
416 .saturating_sub(self.ordinary_bytes);
417 if bytes <= remaining {
418 if let InvokeCommandEvent::Task(TaskEvent::Start { id, .. }) = &event {
419 self.started_task_ids.insert(id.clone());
420 }
421 self.retain_ordinary(event, bytes);
422 return;
423 }
424 if remaining > 0 {
429 if let Some((rebuild, content)) = truncatable_content(event) {
430 let bounded = agent_core::BoundedText::new(&content, remaining);
431 let kept = bounded.retained_bytes;
432 debug_assert!(bounded.truncated);
433 self.counters.truncated_events += 1;
434 self.counters.truncated_payload_bytes += (bytes - kept) as u64;
435 self.retain_ordinary(rebuild(bounded.text), kept);
436 return;
437 }
438 }
440 self.drop_whole(bytes, is_error);
441 }
442
443 fn retain_ordinary(&mut self, event: InvokeCommandEvent, kept_bytes: usize) {
444 self.ordinary_bytes += kept_bytes;
445 self.counters.retained_events += 1;
446 self.counters.retained_payload_bytes += kept_bytes as u64;
447 self.events.push(event);
448 }
449
450 fn retain_control(&mut self, event: InvokeCommandEvent, kept_bytes: usize) {
451 self.control_bytes += kept_bytes;
452 self.counters.retained_events += 1;
453 self.counters.retained_payload_bytes += kept_bytes as u64;
454 self.events.push(event);
455 }
456
457 fn drop_whole(&mut self, bytes: usize, is_error: bool) {
458 self.counters.dropped_events += 1;
459 self.counters.dropped_payload_bytes += bytes as u64;
460 if is_error {
461 self.counters.dropped_error_events += 1;
462 }
463 }
464
465 fn finish(mut self, meter: &InvokeMeter) -> InvokeOutputReport {
466 self.counters.produced_events = meter.produced_events.load(Ordering::SeqCst);
467 self.counters.produced_payload_bytes = meter.produced_payload_bytes.load(Ordering::SeqCst);
468 debug_assert!(
469 self.counters.retained_payload_bytes
470 <= (self.budget.max_retained_payload_bytes as u64)
471 .saturating_add(self.budget.control_reserve_bytes as u64),
472 "retained bytes must never exceed budget + control reserve",
473 );
474 InvokeOutputReport {
475 events: self.events,
476 counters: self.counters,
477 }
478 }
479}
480
481fn is_error_event(event: &InvokeCommandEvent) -> bool {
482 matches!(
483 event,
484 InvokeCommandEvent::Output(CommandOutputEvent::Error { .. })
485 )
486}
487
488#[allow(clippy::type_complexity)]
492fn truncatable_content(
493 event: InvokeCommandEvent,
494) -> Option<(fn(String) -> InvokeCommandEvent, String)> {
495 match event {
496 InvokeCommandEvent::Output(CommandOutputEvent::Text { content }) => Some((
497 |content| InvokeCommandEvent::Output(CommandOutputEvent::Text { content }),
498 content,
499 )),
500 InvokeCommandEvent::Output(CommandOutputEvent::System { content }) => Some((
501 |content| InvokeCommandEvent::Output(CommandOutputEvent::System { content }),
502 content,
503 )),
504 InvokeCommandEvent::Output(CommandOutputEvent::Error { content }) => Some((
505 |content| InvokeCommandEvent::Output(CommandOutputEvent::Error { content }),
506 content,
507 )),
508 _ => None,
509 }
510}
511
512#[cfg(test)]
513mod tests {
514 use super::*;
515
516 fn text(content: &str) -> InvokeCommandEvent {
517 InvokeCommandEvent::Output(CommandOutputEvent::Text {
518 content: content.to_string(),
519 })
520 }
521
522 fn budget(bytes: usize, events: usize, reserve: usize) -> InvokeOutputBudget {
523 InvokeOutputBudget {
524 max_retained_payload_bytes: bytes,
525 max_retained_events: events,
526 control_reserve_bytes: reserve,
527 }
528 }
529
530 fn ingest_all(
531 budget: InvokeOutputBudget,
532 events: Vec<InvokeCommandEvent>,
533 ) -> InvokeOutputReport {
534 let mut state = CollectorState::new(budget);
535 let meter = InvokeMeter::default();
536 for ev in events {
537 meter.produced_events.fetch_add(1, Ordering::SeqCst);
538 meter
539 .produced_payload_bytes
540 .fetch_add(event_payload_bytes(&ev) as u64, Ordering::SeqCst);
541 state.ingest(ev);
542 }
543 state.finish(&meter)
544 }
545
546 fn assert_conservation(report: &InvokeOutputReport) {
547 let c = &report.counters;
548 assert_eq!(
549 c.consumed_payload_bytes,
550 c.retained_payload_bytes + c.truncated_payload_bytes + c.dropped_payload_bytes,
551 "byte conservation"
552 );
553 assert_eq!(
554 c.consumed_events,
555 c.retained_events + c.dropped_events,
556 "event conservation"
557 );
558 assert_eq!(report.events.len() as u64, c.retained_events);
559 }
560
561 #[test]
562 fn under_budget_output_is_retained_verbatim() {
563 let events = vec![
564 text("hello"),
565 text("world"),
566 InvokeCommandEvent::Output(CommandOutputEvent::Done),
567 ];
568 let report = ingest_all(InvokeOutputBudget::default(), events.clone());
569 assert_eq!(report.events, events);
570 assert!(!report.is_limited());
571 assert!(report.limit_notice().is_none());
572 assert!(report.counters.saw_done);
573 assert_eq!(report.counters.retained_payload_bytes, 10);
574 assert_eq!(report.counters.produced_payload_bytes, 10);
575 assert_conservation(&report);
576 }
577
578 #[test]
579 fn head_retention_truncates_boundary_event_then_drops_rest() {
580 let report = ingest_all(
582 budget(10, 100, 0),
583 vec![text("abcdef"), text("ghijkl"), text("mnopqr")],
584 );
585 assert_eq!(report.events, vec![text("abcdef"), text("ghij")]);
586 let c = &report.counters;
587 assert_eq!(c.retained_payload_bytes, 10);
588 assert_eq!(c.truncated_events, 1);
589 assert_eq!(c.truncated_payload_bytes, 2);
590 assert_eq!(c.dropped_events, 1);
591 assert_eq!(c.dropped_payload_bytes, 6);
592 assert_conservation(&report);
593 }
594
595 #[test]
596 fn truncation_is_utf8_char_boundary_safe() {
597 let report = ingest_all(budget(3, 100, 0), vec![text("éé")]);
599 assert_eq!(report.events, vec![text("é")]);
600 assert_eq!(report.counters.retained_payload_bytes, 2);
601 assert_eq!(report.counters.truncated_payload_bytes, 2);
602 assert_conservation(&report);
603 }
604
605 #[test]
606 fn event_budget_caps_tiny_event_floods() {
607 let flood: Vec<_> = (0..500).map(|i| text(&format!("{i}"))).collect();
608 let report = ingest_all(budget(usize::MAX, 16, 0), flood);
609 assert_eq!(report.counters.retained_events, 16);
610 assert_eq!(report.counters.dropped_events, 484);
611 assert_conservation(&report);
612 }
613
614 #[test]
615 fn first_done_marker_survives_budget_exhaustion() {
616 let report = ingest_all(
617 budget(4, 100, 1024),
618 vec![
619 text("aaaa"),
620 text("bbbb"),
621 InvokeCommandEvent::Output(CommandOutputEvent::Done),
622 InvokeCommandEvent::Output(CommandOutputEvent::Done),
623 ],
624 );
625 assert!(report.counters.saw_done);
626 assert_eq!(
627 report.events,
628 vec![
629 text("aaaa"),
630 InvokeCommandEvent::Output(CommandOutputEvent::Done)
631 ]
632 );
633 assert_eq!(report.counters.dropped_events, 2);
635 assert_conservation(&report);
636 }
637
638 #[test]
639 fn task_done_for_retained_start_survives_budget_exhaustion() {
640 let start = InvokeCommandEvent::Task(TaskEvent::Start {
641 id: "t1".into(),
642 label: "L".into(),
643 kind: crate::extensions::tasks::TaskKind::Download,
644 });
645 let done = InvokeCommandEvent::Task(TaskEvent::Done {
646 id: "t1".into(),
647 error: None,
648 });
649 let report = ingest_all(
650 budget(3, 100, 1024),
651 vec![start.clone(), text("xxxxxxxx"), done.clone(), done.clone()],
652 );
653 assert_eq!(report.events, vec![start, done]);
656 assert_eq!(report.counters.dropped_events, 2);
657 assert_conservation(&report);
658 }
659
660 #[test]
661 fn task_done_for_unknown_task_does_not_consume_control_reserve() {
662 let done = InvokeCommandEvent::Task(TaskEvent::Done {
663 id: "ghost".into(),
664 error: None,
665 });
666 let report = ingest_all(budget(0, 0, 1024), vec![done]);
667 assert!(report.events.is_empty());
668 assert_eq!(report.counters.dropped_events, 1);
669 assert_conservation(&report);
670 }
671
672 #[test]
673 fn dropped_error_events_are_counted_and_flagged_in_notice() {
674 let report = ingest_all(
675 budget(2, 100, 0),
676 vec![
677 text("xx"),
678 InvokeCommandEvent::Output(CommandOutputEvent::Error {
679 content: "boom".into(),
680 }),
681 ],
682 );
683 assert_eq!(report.counters.dropped_error_events, 1);
684 let notice = report.limit_notice().expect("limited");
685 assert!(notice.severity_error);
686 assert!(notice.message.contains("including 1 error events"));
687 assert_conservation(&report);
688 }
689
690 #[test]
695 fn disabled_budget_retains_entire_flood() {
696 let flood: Vec<_> = (0..1000).map(|_| text(&"z".repeat(1024))).collect();
697 let report = ingest_all(InvokeOutputBudget::unlimited(), flood);
698 let c = &report.counters;
699 assert_eq!(c.retained_payload_bytes, 1_024_000);
700 assert_eq!(c.retained_events, 1000);
701 assert_eq!(c.dropped_events, 0);
702 assert_eq!(c.truncated_events, 0);
703 assert!(
704 c.retained_payload_bytes > INVOKE_RETAINED_BYTE_BUDGET as u64,
705 "an unlimited budget must visibly violate the production ceiling"
706 );
707 assert_conservation(&report);
708 }
709
710 #[tokio::test]
713 async fn produced_counts_survive_closed_collector() {
714 let (sink, collector) = invoke_event_channel(InvokeOutputBudget::default());
715 sink.send(text("a")).await.expect("open");
716 drop(collector);
718 assert!(matches!(sink.send(text("bb")).await, Err(InvokeSinkClosed)));
719 assert!(sink.is_closed());
720 }
721}