Skip to main content

camel_processor/
aggregator.rs

1use std::collections::HashMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::{Arc, Mutex};
5use std::task::{Context, Poll};
6use std::time::{Duration, Instant};
7
8use tokio::sync::mpsc;
9use tokio::task::JoinHandle;
10use tokio_util::sync::CancellationToken;
11use tower::Service;
12
13use async_trait::async_trait;
14use camel_api::{
15    CamelError, StepLifecycle, StepShutdownReason,
16    aggregator::{
17        AggregationStrategy, AggregatorConfig, CompletionCondition, CompletionMode,
18        CompletionReason, CorrelationStrategy,
19    },
20    body::Body,
21    exchange::Exchange,
22    message::Message,
23};
24use camel_language_api::Language;
25
26pub type SharedLanguageRegistry = Arc<std::sync::Mutex<HashMap<String, Arc<dyn Language>>>>;
27
28pub const CAMEL_AGGREGATOR_PENDING: &str = "CamelAggregatorPending";
29pub const CAMEL_AGGREGATED_SIZE: &str = "CamelAggregatedSize";
30pub const CAMEL_AGGREGATED_KEY: &str = "CamelAggregatedKey";
31pub const CAMEL_AGGREGATED_COMPLETION_REASON: &str = "CamelAggregatedCompletionReason";
32
33/// Internal bucket structure with timestamp tracking for TTL eviction.
34struct Bucket {
35    exchanges: Vec<Exchange>,
36    last_updated: Instant,
37}
38
39impl Bucket {
40    fn new() -> Self {
41        Self {
42            exchanges: Vec::new(),
43            last_updated: Instant::now(),
44        }
45    }
46
47    fn push(&mut self, exchange: Exchange) {
48        self.exchanges.push(exchange);
49        self.last_updated = Instant::now();
50    }
51
52    fn is_expired(&self, ttl: Duration) -> bool {
53        Instant::now().duration_since(self.last_updated) >= ttl
54    }
55}
56
57#[derive(Clone)]
58pub struct AggregatorService {
59    config: AggregatorConfig,
60    buckets: Arc<Mutex<HashMap<String, Bucket>>>,
61    timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
62    timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
63    late_tx: mpsc::Sender<Exchange>,
64    language_registry: SharedLanguageRegistry,
65    route_cancel: CancellationToken,
66    /// Handle to the background TTL-sweep task. `None` when `config.bucket_ttl`
67    /// is `None` (no TTL → no sweep). When the TTL is set, this is populated
68    /// at construction by the auto-spawn in `new` — the caller does not need
69    /// to start the sweep explicitly. The task is bound to `route_cancel` —
70    /// cancelling the route token aborts the task. This is the fix for
71    /// R3-C1's "no background sweep" half: a flood of unique keys within the
72    /// inline-retain window can no longer grow unbounded between calls.
73    sweep_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
74}
75
76impl std::fmt::Debug for AggregatorService {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        f.debug_struct("AggregatorService").finish_non_exhaustive()
79    }
80}
81
82impl Drop for AggregatorService {
83    fn drop(&mut self) {
84        // Defense-in-depth for the R3-C1 sweep: abort the background task when
85        // the service is dropped so it cannot leak even if the route owner
86        // forgets to cancel `route_cancel` (the primary shutdown path).
87        // `abort()` on an already-finished task is a no-op, so this is safe
88        // alongside the `select!` cancel path. This also gives `sweep_handle`
89        // a production reader so it is not dead-code.
90        if let Some(handle) = self
91            .sweep_handle
92            .lock()
93            .unwrap_or_else(|e| e.into_inner())
94            .take()
95        {
96            handle.abort();
97        }
98    }
99}
100
101impl AggregatorService {
102    /// Lifecycle invariant: `route_cancel` is owned by the route. Construction
103    /// is runtime-free (no `tokio::spawn` here) so callers that build an
104    /// `AggregatorService` outside a tokio runtime — e.g. route-spec tests —
105    /// do not panic. When `config.bucket_ttl` is `Some`, the TTL-sweep task is
106    /// spawned LAZILY on the first `poll_ready` (a runtime is guaranteed
107    /// there) and bound to `route_cancel` via `select!`. The route owner MUST
108    /// cancel it on shutdown; `Drop` also aborts it as defense-in-depth.
109    pub fn new(
110        config: AggregatorConfig,
111        late_tx: mpsc::Sender<Exchange>,
112        language_registry: SharedLanguageRegistry,
113        route_cancel: CancellationToken,
114    ) -> Self {
115        // R3-M2: at least one memory-release bound is mandatory.
116        config.validate().expect(
117            // allow-unwrap
118            // fail-closed startup invariant (ADR-0033); a config without a bound is a programmer/operator error.
119            "AggregatorService::new: config failed validation \
120             (need max_buckets, completionTimeout, or bucket_ttl)",
121        );
122
123        // R3-M2 advisory: Size/Predicate-only completion (no Timeout) with no
124        // bucket_ttl means buckets accumulate until max_buckets is hit. Bounded
125        // by max_buckets (validated above) but worth surfacing to operators.
126        let has_timeout = has_timeout_condition(&config.completion);
127        if !has_timeout && config.bucket_ttl.is_none() {
128            tracing::warn!(
129                "Aggregator configured with Size/Predicate completion and no \
130                 bucket_ttl: buckets accumulate until max_buckets is reached"
131            );
132        }
133
134        // Build the shared buckets map up front so the sweep task can
135        // share it via Arc::clone.
136        let buckets: Arc<Mutex<HashMap<String, Bucket>>> = Arc::new(Mutex::new(HashMap::new()));
137
138        // R3-C1 Batch 1: the TTL-sweep is NOT spawned here — construction
139        // must stay runtime-free so callers that build an AggregatorService
140        // outside a tokio runtime (route-spec tests) do not panic on
141        // `tokio::spawn`. The sweep is spawned lazily on the first
142        // `poll_ready`. The inline `guard.retain(...)` per `call` evicts
143        // expired buckets regardless, so the cap + TTL invariants hold
144        // whether or not the sweep has started yet.
145
146        Self {
147            config,
148            buckets,
149            timeout_tasks: Arc::new(Mutex::new(HashMap::new())),
150            timeout_handles: Arc::new(Mutex::new(HashMap::new())),
151            late_tx,
152            language_registry,
153            route_cancel,
154            sweep_handle: Arc::new(Mutex::new(None)),
155        }
156    }
157
158    pub fn config(&self) -> &AggregatorConfig {
159        &self.config
160    }
161
162    pub fn has_timeout(&self) -> bool {
163        has_timeout_condition(&self.config.completion)
164    }
165
166    pub fn force_complete_all(&self) {
167        let mut buckets_guard = self.buckets.lock().unwrap_or_else(|e| e.into_inner());
168        let keys: Vec<String> = buckets_guard.keys().cloned().collect();
169
170        for key in keys {
171            if let Some(bucket) = buckets_guard.remove(&key) {
172                if self.config.force_completion_on_stop {
173                    cancel_timeout_task_with_handle(
174                        &key,
175                        &self.timeout_tasks,
176                        &self.timeout_handles,
177                    );
178                    match aggregate(bucket.exchanges, &self.config.strategy) {
179                        Ok(mut result) => {
180                            result.set_property(
181                                CAMEL_AGGREGATED_COMPLETION_REASON,
182                                serde_json::json!(CompletionReason::Stop.as_str()),
183                            );
184                            if self.late_tx.try_send(result).is_err() {
185                                tracing::warn!(
186                                    key = %key,
187                                    "aggregator force-complete emit dropped: late channel full"
188                                );
189                            }
190                        }
191                        Err(e) => {
192                            // log-policy: handler-owned
193                            tracing::warn!(
194                                key = %key,
195                                error = %e,
196                                "aggregation failed in force_complete_all"
197                            );
198                        }
199                    }
200                } else {
201                    cancel_timeout_task_with_handle(
202                        &key,
203                        &self.timeout_tasks,
204                        &self.timeout_handles,
205                    );
206                }
207            }
208        }
209    }
210
211    /// Graceful shutdown: cancel all outstanding timeout tasks and await their
212    /// JoinHandles (with a 5s deadline) so that no tasks are leaked.
213    pub(crate) async fn shutdown_inner(&self) {
214        // Cancel all timeout cancellation tokens.
215        {
216            let mut guard = self.timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
217            for token in guard.values() {
218                token.cancel();
219            }
220            guard.clear();
221        };
222
223        // Remove and collect all JoinHandles.
224        let handles: Vec<JoinHandle<()>> = {
225            let mut guard = self
226                .timeout_handles
227                .lock()
228                .unwrap_or_else(|e| e.into_inner());
229            guard.drain().map(|(_, handle)| handle).collect()
230        };
231
232        if handles.is_empty() {
233            return;
234        }
235
236        // Await all handles with a deadline.
237        let _ = tokio::time::timeout(Duration::from_secs(5), async {
238            for handle in handles {
239                let _ = handle.await;
240            }
241        })
242        .await;
243    }
244}
245
246#[async_trait]
247impl StepLifecycle for AggregatorService {
248    fn name(&self) -> &'static str {
249        "aggregator"
250    }
251
252    async fn shutdown(&self, reason: StepShutdownReason) -> Result<(), CamelError> {
253        tracing::debug!(reason = ?reason, "Aggregator shutdown via StepLifecycle");
254        self.shutdown_inner().await;
255        Ok(())
256    }
257}
258
259pub fn has_timeout_condition(mode: &CompletionMode) -> bool {
260    match mode {
261        CompletionMode::Single(CompletionCondition::Timeout(_)) => true,
262        CompletionMode::Any(conditions) => conditions
263            .iter()
264            .any(|c| matches!(c, CompletionCondition::Timeout(_))),
265        _ => false,
266    }
267}
268
269impl Service<Exchange> for AggregatorService {
270    type Response = Exchange;
271    type Error = CamelError;
272    type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
273
274    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), CamelError>> {
275        // R3-C1: lazily spawn the TTL-sweep on the first readiness poll. A
276        // tokio runtime is guaranteed here (the runtime driving the service),
277        // unlike at construction. Single-spawn via the lock + is_none check.
278        if let Some(ttl) = self.config.bucket_ttl {
279            let mut g = self.sweep_handle.lock().unwrap_or_else(|e| e.into_inner());
280            if g.is_none() {
281                let interval = std::cmp::max(ttl / 2, Duration::from_millis(50));
282                let buckets = Arc::clone(&self.buckets);
283                let cancel = self.route_cancel.clone();
284                *g = Some(tokio::spawn(async move {
285                    loop {
286                        tokio::select! {
287                            _ = cancel.cancelled() => break,
288                            _ = tokio::time::sleep(interval) => {
289                                let mut guard =
290                                    buckets.lock().unwrap_or_else(|e| e.into_inner());
291                                guard.retain(|_, b| !b.is_expired(ttl));
292                            }
293                        }
294                    }
295                }));
296            }
297        }
298        Poll::Ready(Ok(()))
299    }
300
301    fn call(&mut self, exchange: Exchange) -> Self::Future {
302        let config = self.config.clone();
303        let buckets = Arc::clone(&self.buckets);
304        let timeout_tasks = Arc::clone(&self.timeout_tasks);
305        let timeout_handles = Arc::clone(&self.timeout_handles);
306        let late_tx = self.late_tx.clone();
307        let language_registry = Arc::clone(&self.language_registry);
308        let route_cancel = self.route_cancel.clone();
309
310        Box::pin(async move {
311            let key_value =
312                extract_correlation_key(&exchange, &config.correlation, &language_registry).await?;
313
314            let key_str = serde_json::to_string(&key_value)
315                .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
316
317            let completed_bucket = {
318                let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
319
320                if let Some(ttl) = config.bucket_ttl {
321                    guard.retain(|_, bucket| !bucket.is_expired(ttl));
322                }
323
324                if let Some(max) = config.max_buckets
325                    && !guard.contains_key(&key_str)
326                    && guard.len() >= max
327                {
328                    tracing::warn!(
329                        max_buckets = max,
330                        correlation_key = %key_str,
331                        "Aggregator reached max buckets limit, rejecting new correlation key"
332                    );
333                    return Err(CamelError::ProcessorError(format!(
334                        "Aggregator reached maximum {} buckets",
335                        max
336                    )));
337                }
338
339                let bucket = guard.entry(key_str.clone()).or_insert_with(Bucket::new);
340                bucket.push(exchange);
341
342                let (is_complete, reason) =
343                    check_sync_completion(&config.completion, &bucket.exchanges);
344
345                if is_complete {
346                    let exchanges = guard.remove(&key_str).map(|b| b.exchanges);
347                    (exchanges, reason)
348                } else {
349                    (None, CompletionReason::Size) // placeholder; reason unused when None
350                }
351            };
352
353            if completed_bucket.0.is_none() && has_timeout_condition(&config.completion) {
354                let timeout_dur = extract_timeout_duration(&config.completion);
355                if let Some(timeout) = timeout_dur {
356                    // R3-M3: bound the number of concurrently-live per-bucket
357                    // timeout tasks. When the cap is reached, skip the dedicated
358                    // spawn — the bucket relies on bucket_ttl eviction (graceful
359                    // degradation). max_buckets already caps total buckets, so
360                    // memory stays bounded regardless.
361                    let live_count = timeout_handles
362                        .lock()
363                        .unwrap_or_else(|e| e.into_inner())
364                        .len();
365                    if live_count >= config.max_timeout_tasks {
366                        tracing::warn!(
367                            live_timeout_tasks = live_count,
368                            max_timeout_tasks = config.max_timeout_tasks,
369                            correlation_key = %key_str,
370                            "Aggregator timeout-task cap reached; bucket will rely on \
371                             bucket_ttl eviction instead of a dedicated timeout task"
372                        );
373                    } else {
374                        // Cancel old token for this key (if any).
375                        {
376                            let tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
377                            if let Some(existing) = tt_guard.get(&key_str) {
378                                existing.cancel();
379                            }
380                        }
381                        // Remove old handle if present.
382                        {
383                            let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
384                            if let Some(old) = hh.remove(&key_str) {
385                                old.abort();
386                            }
387                        }
388                        let cancel = CancellationToken::new();
389                        timeout_tasks
390                            .lock()
391                            .unwrap_or_else(|e| e.into_inner())
392                            .insert(key_str.clone(), cancel.clone());
393                        let handle = spawn_timeout_task(
394                            key_str.clone(),
395                            timeout,
396                            cancel,
397                            buckets.clone(),
398                            timeout_tasks.clone(),
399                            timeout_handles.clone(),
400                            late_tx,
401                            config.strategy.clone(),
402                            config.discard_on_timeout,
403                            route_cancel,
404                        );
405                        timeout_handles
406                            .lock()
407                            .unwrap_or_else(|e| e.into_inner())
408                            .insert(key_str.clone(), handle);
409                    }
410                }
411            }
412
413            if let Some(exchanges) = completed_bucket.0 {
414                cancel_timeout_task_with_handle(&key_str, &timeout_tasks, &timeout_handles);
415                let reason = completed_bucket.1;
416                let size = exchanges.len();
417                let mut result = aggregate(exchanges, &config.strategy)?;
418                result.set_property(CAMEL_AGGREGATED_SIZE, serde_json::json!(size as u64));
419                result.set_property(CAMEL_AGGREGATED_KEY, key_value);
420                result.set_property(
421                    CAMEL_AGGREGATED_COMPLETION_REASON,
422                    serde_json::json!(reason.as_str()),
423                );
424                Ok(result)
425            } else {
426                let mut pending = Exchange::new(Message {
427                    headers: Default::default(),
428                    body: Body::Empty,
429                });
430                pending.set_property(CAMEL_AGGREGATOR_PENDING, serde_json::json!(true));
431                Ok(pending)
432            }
433        })
434    }
435}
436
437async fn extract_correlation_key(
438    exchange: &Exchange,
439    strategy: &CorrelationStrategy,
440    registry: &SharedLanguageRegistry,
441) -> Result<serde_json::Value, CamelError> {
442    match strategy {
443        CorrelationStrategy::HeaderName(h) => {
444            exchange.input.headers.get(h).cloned().ok_or_else(|| {
445                CamelError::ProcessorError(format!(
446                    "Aggregator: missing correlation key header '{}'",
447                    h
448                ))
449            })
450        }
451        CorrelationStrategy::Expression { expr, language } => {
452            let expression = {
453                let reg = registry.lock().unwrap_or_else(|e| e.into_inner());
454                let lang = reg.get(language).ok_or_else(|| {
455                    CamelError::ProcessorError(format!(
456                        "Aggregator: language '{}' not found in registry",
457                        language
458                    ))
459                })?;
460                lang.create_expression(expr)
461                    .map_err(|e| CamelError::ProcessorError(e.to_string()))?
462            };
463            let value = expression
464                .evaluate(exchange)
465                .await
466                .map_err(|e| CamelError::ProcessorError(e.to_string()))?;
467            if value.is_null() {
468                return Err(CamelError::ProcessorError(format!(
469                    "Aggregator: correlation expression '{}' evaluated to null",
470                    expr
471                )));
472            }
473            Ok(value)
474        }
475        CorrelationStrategy::Fn(f) => f(exchange).map(serde_json::Value::String).ok_or_else(|| {
476            CamelError::ProcessorError("Aggregator: correlation function returned None".to_string())
477        }),
478    }
479}
480
481fn check_sync_completion(
482    mode: &CompletionMode,
483    exchanges: &[Exchange],
484) -> (bool, CompletionReason) {
485    match mode {
486        CompletionMode::Single(cond) => check_single(cond, exchanges),
487        CompletionMode::Any(conditions) => {
488            for cond in conditions {
489                if let CompletionCondition::Timeout(_) = cond {
490                    continue;
491                }
492                let (done, reason) = check_single(cond, exchanges);
493                if done {
494                    return (true, reason);
495                }
496            }
497            (false, CompletionReason::Size)
498        }
499    }
500}
501
502fn check_single(cond: &CompletionCondition, exchanges: &[Exchange]) -> (bool, CompletionReason) {
503    match cond {
504        CompletionCondition::Size(n) => (exchanges.len() >= *n, CompletionReason::Size),
505        CompletionCondition::Predicate(pred) => (pred(exchanges), CompletionReason::Predicate),
506        CompletionCondition::Timeout(_) => (false, CompletionReason::Timeout),
507    }
508}
509
510fn extract_timeout_duration(mode: &CompletionMode) -> Option<Duration> {
511    match mode {
512        CompletionMode::Single(CompletionCondition::Timeout(d)) => Some(*d),
513        CompletionMode::Any(conditions) => conditions.iter().find_map(|c| {
514            if let CompletionCondition::Timeout(d) = c {
515                Some(*d)
516            } else {
517                None
518            }
519        }),
520        _ => None,
521    }
522}
523
524fn cancel_timeout_task(key: &str, timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>) {
525    let mut guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
526    if let Some(token) = guard.remove(key) {
527        token.cancel();
528    }
529}
530
531/// Also removes the stored JoinHandle for a cancelled/completed timeout task.
532fn cancel_timeout_task_with_handle(
533    key: &str,
534    timeout_tasks: &Arc<Mutex<HashMap<String, CancellationToken>>>,
535    timeout_handles: &Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
536) {
537    cancel_timeout_task(key, timeout_tasks);
538    let mut guard = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
539    guard.remove(key);
540}
541
542#[allow(clippy::too_many_arguments)]
543fn spawn_timeout_task(
544    key: String,
545    timeout: Duration,
546    cancel: CancellationToken,
547    buckets: Arc<Mutex<HashMap<String, Bucket>>>,
548    timeout_tasks: Arc<Mutex<HashMap<String, CancellationToken>>>,
549    timeout_handles: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
550    late_tx: mpsc::Sender<Exchange>,
551    strategy: AggregationStrategy,
552    discard: bool,
553    _route_cancel: CancellationToken,
554) -> JoinHandle<()> {
555    let cancel_clone = cancel.clone();
556    tokio::spawn(async move {
557        tokio::select! {
558            _ = tokio::time::sleep(timeout) => {
559                let should_proceed = {
560                    let mut tt_guard = timeout_tasks.lock().unwrap_or_else(|e| e.into_inner());
561                    if cancel_clone.is_cancelled() {
562                        false
563                    } else {
564                        tt_guard.remove(&key);
565                        true
566                    }
567                };
568                if !should_proceed {
569                    return;
570                }
571                // Clean up our own JoinHandle from the map on natural completion.
572                // Without this, the handle entry leaks until route shutdown.
573                {
574                    let mut hh = timeout_handles.lock().unwrap_or_else(|e| e.into_inner());
575                    hh.remove(&key);
576                }
577                let bucket_exchanges = {
578                    let mut guard = buckets.lock().unwrap_or_else(|e| e.into_inner());
579                    guard.remove(&key).map(|b| b.exchanges)
580                };
581                if let Some(exchanges) = bucket_exchanges
582                    && !discard
583                {
584                    match aggregate(exchanges, &strategy) {
585                        Ok(mut result) => {
586                            result.set_property(
587                                CAMEL_AGGREGATED_COMPLETION_REASON,
588                                serde_json::json!(CompletionReason::Timeout.as_str()),
589                            );
590                            if late_tx.try_send(result).is_err() {
591                                tracing::warn!(
592                                    key = %key,
593                                    "aggregator timeout emit dropped: late channel full"
594                                );
595                            }
596                        }
597                        Err(e) => {
598                            // log-policy: handler-owned
599                            tracing::warn!(
600                                key = %key,
601                                error = %e,
602                                "aggregation failed in timeout task"
603                            );
604                        }
605                    }
606                }
607            }
608            _ = cancel_clone.cancelled() => {}
609        }
610    })
611}
612
613fn aggregate(
614    exchanges: Vec<Exchange>,
615    strategy: &AggregationStrategy,
616) -> Result<Exchange, CamelError> {
617    match strategy {
618        AggregationStrategy::CollectAll => {
619            let bodies: Vec<serde_json::Value> = exchanges
620                .into_iter()
621                .map(|e| match e.input.body {
622                    Body::Json(v) => v,
623                    Body::Text(s) => serde_json::Value::String(s),
624                    Body::Xml(s) => serde_json::Value::String(s),
625                    Body::Bytes(b) => {
626                        serde_json::Value::String(String::from_utf8_lossy(&b).into_owned())
627                    }
628                    Body::Empty => serde_json::Value::Null,
629                    Body::Stream(s) => serde_json::json!({
630                        "_stream": {
631                            "origin": s.metadata.origin,
632                            "placeholder": true,
633                            "hint": "Materialize exchange body with .into_bytes() before aggregation if content needed"
634                        }
635                    }),
636                })
637                .collect();
638            Ok(Exchange::new(Message {
639                headers: Default::default(),
640                body: Body::Json(serde_json::Value::Array(bodies)),
641            }))
642        }
643        AggregationStrategy::Custom(f) => {
644            let mut iter = exchanges.into_iter();
645            let first = iter.next().ok_or_else(|| {
646                CamelError::ProcessorError("Aggregator: empty bucket".to_string())
647            })?;
648            Ok(iter.fold(first, |acc, next| f(acc, next)))
649        }
650    }
651}
652
653#[cfg(test)]
654mod tests {
655    use super::*;
656    use std::collections::HashMap;
657
658    use camel_api::{
659        StepLifecycle, StepShutdownReason,
660        aggregator::{AggregationStrategy, AggregatorConfig},
661        body::Body,
662        exchange::Exchange,
663        message::Message,
664    };
665    use tokio::sync::mpsc;
666    use tokio_util::sync::CancellationToken;
667    use tower::ServiceExt;
668
669    fn make_exchange(header: &str, value: &str, body: &str) -> Exchange {
670        let mut msg = Message {
671            headers: Default::default(),
672            body: Body::Text(body.to_string()),
673        };
674        msg.headers
675            .insert(header.to_string(), serde_json::json!(value));
676        Exchange::new(msg)
677    }
678
679    fn config_size(n: usize) -> AggregatorConfig {
680        AggregatorConfig::correlate_by("orderId")
681            .complete_when_size(n)
682            .build()
683            .unwrap()
684    }
685
686    fn new_test_svc(config: AggregatorConfig) -> AggregatorService {
687        let (tx, _rx) = mpsc::channel(256);
688        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
689        let cancel = CancellationToken::new();
690        AggregatorService::new(config, tx, registry, cancel)
691    }
692
693    #[tokio::test]
694    async fn test_pending_exchange_not_yet_complete() {
695        let mut svc = new_test_svc(config_size(3));
696        let ex = make_exchange("orderId", "A", "first");
697        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
698        assert!(matches!(result.input.body, Body::Empty));
699        assert_eq!(
700            result.property(CAMEL_AGGREGATOR_PENDING),
701            Some(&serde_json::json!(true))
702        );
703    }
704
705    #[tokio::test]
706    async fn test_completes_on_size() {
707        let mut svc = new_test_svc(config_size(3));
708        for _ in 0..2 {
709            let ex = make_exchange("orderId", "A", "item");
710            let r = svc.ready().await.unwrap().call(ex).await.unwrap();
711            assert!(matches!(r.input.body, Body::Empty));
712        }
713        let ex = make_exchange("orderId", "A", "last");
714        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
715        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
716        assert_eq!(
717            result.property(CAMEL_AGGREGATED_SIZE),
718            Some(&serde_json::json!(3u64))
719        );
720    }
721
722    #[tokio::test]
723    async fn test_collect_all_produces_json_array() {
724        let mut svc = new_test_svc(config_size(2));
725        svc.ready()
726            .await
727            .unwrap()
728            .call(make_exchange("orderId", "A", "alpha"))
729            .await
730            .unwrap();
731        let result = svc
732            .ready()
733            .await
734            .unwrap()
735            .call(make_exchange("orderId", "A", "beta"))
736            .await
737            .unwrap();
738        let Body::Json(v) = &result.input.body else {
739            panic!("expected Body::Json")
740        };
741        let arr = v.as_array().unwrap();
742        assert_eq!(arr.len(), 2);
743        assert_eq!(arr[0], serde_json::json!("alpha"));
744        assert_eq!(arr[1], serde_json::json!("beta"));
745    }
746
747    #[tokio::test]
748    async fn test_two_keys_independent_buckets() {
749        // completionSize=3 so we can test that A and B accumulate independently.
750        let mut svc = new_test_svc(config_size(3));
751        svc.ready()
752            .await
753            .unwrap()
754            .call(make_exchange("orderId", "A", "a1"))
755            .await
756            .unwrap();
757        svc.ready()
758            .await
759            .unwrap()
760            .call(make_exchange("orderId", "B", "b1"))
761            .await
762            .unwrap();
763        svc.ready()
764            .await
765            .unwrap()
766            .call(make_exchange("orderId", "A", "a2"))
767            .await
768            .unwrap();
769        // A has 2 items, B has 1 item — neither complete yet
770        let ra = svc
771            .ready()
772            .await
773            .unwrap()
774            .call(make_exchange("orderId", "A", "a3"))
775            .await
776            .unwrap();
777        // A now has 3 → completes
778        assert!(matches!(ra.input.body, Body::Json(_)));
779        // B only has 1 → still pending
780        let rb = svc
781            .ready()
782            .await
783            .unwrap()
784            .call(make_exchange("orderId", "B", "b_check"))
785            .await
786            .unwrap();
787        assert!(matches!(rb.input.body, Body::Empty));
788    }
789
790    #[tokio::test]
791    async fn test_bucket_resets_after_completion() {
792        let mut svc = new_test_svc(config_size(2));
793        svc.ready()
794            .await
795            .unwrap()
796            .call(make_exchange("orderId", "A", "x"))
797            .await
798            .unwrap();
799        svc.ready()
800            .await
801            .unwrap()
802            .call(make_exchange("orderId", "A", "x"))
803            .await
804            .unwrap(); // completes
805        // New bucket starts
806        let r = svc
807            .ready()
808            .await
809            .unwrap()
810            .call(make_exchange("orderId", "A", "new"))
811            .await
812            .unwrap();
813        assert!(matches!(r.input.body, Body::Empty)); // pending again
814    }
815
816    #[tokio::test]
817    async fn test_completion_size_1_emits_immediately() {
818        let mut svc = new_test_svc(config_size(1));
819        let ex = make_exchange("orderId", "A", "solo");
820        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
821        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
822    }
823
824    #[tokio::test]
825    async fn test_custom_aggregation_strategy() {
826        use camel_api::aggregator::AggregationFn;
827        use std::sync::Arc;
828
829        let f: AggregationFn = Arc::new(|mut acc: Exchange, next: Exchange| {
830            let combined = format!(
831                "{}+{}",
832                acc.input.body.as_text().unwrap_or(""),
833                next.input.body.as_text().unwrap_or("")
834            );
835            acc.input.body = Body::Text(combined);
836            acc
837        });
838        let config = AggregatorConfig::correlate_by("key")
839            .complete_when_size(2)
840            .strategy(AggregationStrategy::Custom(f))
841            .build()
842            .unwrap();
843        let mut svc = new_test_svc(config);
844        svc.ready()
845            .await
846            .unwrap()
847            .call(make_exchange("key", "X", "hello"))
848            .await
849            .unwrap();
850        let result = svc
851            .ready()
852            .await
853            .unwrap()
854            .call(make_exchange("key", "X", "world"))
855            .await
856            .unwrap();
857        assert_eq!(result.input.body.as_text(), Some("hello+world"));
858    }
859
860    #[tokio::test]
861    async fn test_completion_predicate() {
862        let config = AggregatorConfig::correlate_by("key")
863            .complete_when(|bucket| {
864                bucket
865                    .iter()
866                    .any(|e| e.input.body.as_text() == Some("DONE"))
867            })
868            .build()
869            .unwrap();
870        let mut svc = new_test_svc(config);
871        svc.ready()
872            .await
873            .unwrap()
874            .call(make_exchange("key", "K", "first"))
875            .await
876            .unwrap();
877        svc.ready()
878            .await
879            .unwrap()
880            .call(make_exchange("key", "K", "second"))
881            .await
882            .unwrap();
883        let result = svc
884            .ready()
885            .await
886            .unwrap()
887            .call(make_exchange("key", "K", "DONE"))
888            .await
889            .unwrap();
890        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
891    }
892
893    #[tokio::test]
894    async fn test_missing_header_returns_error() {
895        let mut svc = new_test_svc(config_size(2));
896        let msg = Message {
897            headers: Default::default(),
898            body: Body::Text("no key".into()),
899        };
900        let ex = Exchange::new(msg);
901        let result = svc.ready().await.unwrap().call(ex).await;
902        assert!(result.is_err());
903        assert!(matches!(
904            result.unwrap_err(),
905            camel_api::CamelError::ProcessorError(_)
906        ));
907    }
908
909    #[tokio::test]
910    async fn test_cloned_service_shares_state() {
911        let svc1 = new_test_svc(config_size(2));
912        let mut svc2 = svc1.clone();
913        // send first exchange via svc1
914        svc1.clone()
915            .ready()
916            .await
917            .unwrap()
918            .call(make_exchange("orderId", "A", "from-svc1"))
919            .await
920            .unwrap();
921        // send second exchange via svc2 — should complete because same Arc<Mutex>
922        let result = svc2
923            .ready()
924            .await
925            .unwrap()
926            .call(make_exchange("orderId", "A", "from-svc2"))
927            .await
928            .unwrap();
929        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
930    }
931
932    #[tokio::test]
933    async fn test_camel_aggregated_key_property_set() {
934        let mut svc = new_test_svc(config_size(1));
935        let ex = make_exchange("orderId", "ORDER-42", "body");
936        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
937        assert_eq!(
938            result.property(CAMEL_AGGREGATED_KEY),
939            Some(&serde_json::json!("ORDER-42"))
940        );
941    }
942
943    #[tokio::test]
944    async fn test_aggregator_enforces_max_buckets() {
945        let config = AggregatorConfig::correlate_by("orderId")
946            .complete_when_size(2)
947            .max_buckets(3)
948            .build()
949            .unwrap();
950
951        let mut svc = new_test_svc(config);
952
953        // Create 3 different correlation keys (fills limit)
954        for i in 0..3 {
955            let ex = make_exchange("orderId", &format!("key-{}", i), "body");
956            let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
957        }
958
959        // 4th key should be rejected
960        let ex = make_exchange("orderId", "key-4", "body");
961        let result = svc.ready().await.unwrap().call(ex).await;
962
963        assert!(result.is_err(), "Should reject when max buckets reached");
964        let err = result.unwrap_err().to_string();
965        assert!(
966            err.contains("maximum"),
967            "Error message should contain 'maximum': {}",
968            err
969        );
970    }
971
972    #[tokio::test]
973    async fn test_max_buckets_allows_existing_key() {
974        let config = AggregatorConfig::correlate_by("orderId")
975            .complete_when_size(5) // Large size so bucket doesn't complete
976            .max_buckets(2)
977            .build()
978            .unwrap();
979
980        let mut svc = new_test_svc(config);
981
982        // Create 2 different correlation keys (fills limit)
983        let ex1 = make_exchange("orderId", "key-A", "body1");
984        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
985        let ex2 = make_exchange("orderId", "key-B", "body2");
986        let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
987
988        // Should still allow adding to existing key
989        let ex3 = make_exchange("orderId", "key-A", "body3");
990        let result = svc.ready().await.unwrap().call(ex3).await;
991        assert!(
992            result.is_ok(),
993            "Should allow adding to existing bucket even at max limit"
994        );
995    }
996
997    #[tokio::test]
998    async fn test_bucket_ttl_eviction() {
999        let config = AggregatorConfig::correlate_by("orderId")
1000            .complete_when_size(10) // Large size so bucket doesn't complete normally
1001            .bucket_ttl(Duration::from_millis(50))
1002            .build()
1003            .unwrap();
1004
1005        let mut svc = new_test_svc(config);
1006
1007        // Create a bucket
1008        let ex1 = make_exchange("orderId", "key-A", "body1");
1009        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1010
1011        // Wait for TTL to expire
1012        tokio::time::sleep(Duration::from_millis(100)).await;
1013
1014        // Create a new bucket - this should trigger eviction of the old one
1015        let ex2 = make_exchange("orderId", "key-B", "body2");
1016        let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1017
1018        // The expired bucket should have been evicted, so we should be able to
1019        // add a new key-A bucket again
1020        let ex3 = make_exchange("orderId", "key-A", "body3");
1021        let result = svc.ready().await.unwrap().call(ex3).await;
1022        assert!(result.is_ok(), "Should be able to recreate evicted bucket");
1023    }
1024
1025    #[tokio::test(start_paused = true)]
1026    async fn test_timeout_completes_bucket() {
1027        let config = AggregatorConfig::correlate_by("key")
1028            .complete_on_timeout(Duration::from_millis(100))
1029            .build()
1030            .unwrap();
1031        let mut svc = new_test_svc(config);
1032        let ex = make_exchange("key", "A", "data");
1033        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1034        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_some());
1035
1036        tokio::time::sleep(Duration::from_millis(200)).await;
1037
1038        assert_eq!(
1039            svc.buckets.lock().unwrap().len(),
1040            0,
1041            "bucket should be removed after timeout"
1042        );
1043    }
1044
1045    #[tokio::test(start_paused = true)]
1046    async fn test_timeout_resets_on_new_exchange() {
1047        let config = AggregatorConfig::correlate_by("key")
1048            .complete_on_timeout(Duration::from_millis(150))
1049            .build()
1050            .unwrap();
1051        let mut svc = new_test_svc(config);
1052
1053        let ex1 = make_exchange("key", "A", "first");
1054        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1055
1056        tokio::time::sleep(Duration::from_millis(100)).await;
1057
1058        let ex2 = make_exchange("key", "A", "second");
1059        let _ = svc.ready().await.unwrap().call(ex2).await.unwrap();
1060
1061        tokio::time::sleep(Duration::from_millis(100)).await;
1062
1063        assert_eq!(
1064            svc.buckets.lock().unwrap().len(),
1065            1,
1066            "bucket should still exist — timeout was reset"
1067        );
1068
1069        tokio::time::sleep(Duration::from_millis(100)).await;
1070
1071        assert_eq!(
1072            svc.buckets.lock().unwrap().len(),
1073            0,
1074            "bucket should be gone after timeout fires"
1075        );
1076    }
1077
1078    #[tokio::test]
1079    async fn test_composable_size_and_timeout() {
1080        let config = AggregatorConfig::correlate_by("key")
1081            .complete_on_size_or_timeout(2, Duration::from_millis(200))
1082            .build()
1083            .unwrap();
1084        let mut svc = new_test_svc(config);
1085
1086        let ex1 = make_exchange("key", "A", "first");
1087        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1088        assert!(svc.buckets.lock().unwrap().contains_key("\"A\""));
1089
1090        let ex2 = make_exchange("key", "A", "second");
1091        let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1092        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1093        assert_eq!(
1094            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1095            Some(&serde_json::json!("size"))
1096        );
1097    }
1098
1099    #[tokio::test(start_paused = true)]
1100    async fn test_discard_on_timeout() {
1101        let config = AggregatorConfig::correlate_by("key")
1102            .complete_on_timeout(Duration::from_millis(50))
1103            .discard_on_timeout(true)
1104            .build()
1105            .unwrap();
1106        let (tx, mut rx) = mpsc::channel(256);
1107        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1108        let cancel = CancellationToken::new();
1109        let mut svc = AggregatorService::new(config, tx, registry, cancel);
1110
1111        let ex = make_exchange("key", "A", "data");
1112        let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1113
1114        tokio::time::sleep(Duration::from_millis(100)).await;
1115
1116        assert!(
1117            rx.try_recv().is_err(),
1118            "no emit expected with discard_on_timeout"
1119        );
1120        assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1121        assert!(
1122            svc.timeout_tasks.lock().unwrap().is_empty(),
1123            "timeout task should be cleaned up"
1124        );
1125    }
1126
1127    #[tokio::test]
1128    async fn test_force_completion_on_stop() {
1129        let config = AggregatorConfig::correlate_by("key")
1130            .complete_when_size(10)
1131            .force_completion_on_stop(true)
1132            .build()
1133            .unwrap();
1134        let (tx, mut rx) = mpsc::channel(256);
1135        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1136        let cancel = CancellationToken::new();
1137        let svc = AggregatorService::new(config, tx, registry, cancel);
1138
1139        let mut call_svc = svc.clone();
1140        let ex = make_exchange("key", "A", "data");
1141        let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1142
1143        svc.force_complete_all();
1144
1145        let result = rx.try_recv().expect("should emit on force-complete");
1146        assert!(
1147            result.input.body.as_text().is_some() || matches!(result.input.body, Body::Json(_))
1148        );
1149        assert_eq!(
1150            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1151            Some(&serde_json::json!("stop"))
1152        );
1153    }
1154
1155    #[tokio::test]
1156    async fn test_completion_reason_property_size() {
1157        let config = AggregatorConfig::correlate_by("key")
1158            .complete_when_size(1)
1159            .build()
1160            .unwrap();
1161        let mut svc = new_test_svc(config);
1162        let ex = make_exchange("key", "X", "body");
1163        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1164        assert_eq!(
1165            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1166            Some(&serde_json::json!("size"))
1167        );
1168    }
1169
1170    #[tokio::test]
1171    async fn test_completion_reason_property_predicate() {
1172        let config = AggregatorConfig::correlate_by("key")
1173            .complete_when(|_| true)
1174            .build()
1175            .unwrap();
1176        let mut svc = new_test_svc(config);
1177        let ex = make_exchange("key", "X", "body");
1178        let result = svc.ready().await.unwrap().call(ex).await.unwrap();
1179        assert_eq!(
1180            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1181            Some(&serde_json::json!("predicate"))
1182        );
1183    }
1184
1185    #[tokio::test(start_paused = true)]
1186    async fn test_size_completes_before_timeout() {
1187        let config = AggregatorConfig::correlate_by("key")
1188            .complete_on_size_or_timeout(2, Duration::from_millis(200))
1189            .build()
1190            .unwrap();
1191        let mut svc = new_test_svc(config);
1192
1193        let ex1 = make_exchange("key", "A", "first");
1194        let _ = svc.ready().await.unwrap().call(ex1).await.unwrap();
1195
1196        let ex2 = make_exchange("key", "A", "second");
1197        let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1198
1199        assert!(result.property(CAMEL_AGGREGATOR_PENDING).is_none());
1200        assert_eq!(
1201            result.property(CAMEL_AGGREGATED_COMPLETION_REASON),
1202            Some(&serde_json::json!("size"))
1203        );
1204        assert_eq!(svc.buckets.lock().unwrap().len(), 0);
1205
1206        tokio::time::sleep(Duration::from_millis(300)).await;
1207        assert_eq!(
1208            svc.buckets.lock().unwrap().len(),
1209            0,
1210            "no re-fire after timeout"
1211        );
1212    }
1213
1214    #[tokio::test(start_paused = true)]
1215    async fn test_concurrent_timeout_fire_and_new_exchange() {
1216        let config = AggregatorConfig::correlate_by("key")
1217            .complete_on_size_or_timeout(2, Duration::from_millis(100))
1218            .build()
1219            .unwrap();
1220        let (tx, mut rx) = mpsc::channel(256);
1221        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1222        let cancel = CancellationToken::new();
1223        let mut svc = AggregatorService::new(config, tx, registry, cancel);
1224
1225        let ex = make_exchange("key", "A", "data");
1226        let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1227
1228        // Advance time past timeout — timeout task fires and removes bucket
1229        tokio::time::sleep(Duration::from_millis(150)).await;
1230
1231        // New exchange arrives after timeout — starts a fresh bucket
1232        let ex2 = make_exchange("key", "A", "data2");
1233        let result = svc.ready().await.unwrap().call(ex2).await.unwrap();
1234        assert!(
1235            result.property(CAMEL_AGGREGATOR_PENDING).is_some(),
1236            "should be pending in new bucket"
1237        );
1238
1239        // Drain late emits from timeout
1240        let mut late_count = 0;
1241        while rx.try_recv().is_ok() {
1242            late_count += 1;
1243        }
1244        assert_eq!(
1245            late_count, 1,
1246            "exactly 1 late emit from the timed-out bucket"
1247        );
1248    }
1249
1250    #[tokio::test(start_paused = true)]
1251    async fn test_late_channel_full_drops_with_warning() {
1252        let config = AggregatorConfig::correlate_by("key")
1253            .complete_on_timeout(Duration::from_millis(50))
1254            .build()
1255            .unwrap();
1256        let (tx, mut rx) = mpsc::channel(1);
1257        rx.close();
1258        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1259        let cancel = CancellationToken::new();
1260        let mut svc = AggregatorService::new(config, tx, registry, cancel);
1261
1262        let ex = make_exchange("key", "A", "data");
1263        let _ = svc.ready().await.unwrap().call(ex).await.unwrap();
1264
1265        tokio::time::sleep(Duration::from_millis(100)).await;
1266        assert_eq!(
1267            svc.buckets.lock().unwrap().len(),
1268            0,
1269            "bucket removed despite channel closed"
1270        );
1271    }
1272
1273    #[tokio::test]
1274    async fn test_aggregate_stream_bodies_creates_valid_json() {
1275        use bytes::Bytes;
1276        use camel_api::{Body, StreamBody, StreamMetadata};
1277        use futures::stream;
1278        use tokio::sync::Mutex;
1279
1280        let chunks = vec![Ok(Bytes::from("test"))];
1281        let stream_body = StreamBody {
1282            stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1283            metadata: StreamMetadata {
1284                origin: Some("file:///test.txt".to_string()),
1285                ..Default::default()
1286            },
1287        };
1288
1289        let ex1 = Exchange::new(Message {
1290            headers: Default::default(),
1291            body: Body::Stream(stream_body),
1292        });
1293
1294        let exchanges = vec![ex1];
1295        let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1296
1297        let exchange = result.expect("Expected Ok result");
1298        assert!(
1299            matches!(exchange.input.body, Body::Json(_)),
1300            "Expected Json body"
1301        );
1302
1303        if let Body::Json(value) = exchange.input.body {
1304            let json_str = serde_json::to_string(&value).unwrap();
1305            let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1306
1307            assert!(parsed.is_array(), "Result should be an array");
1308            let arr = parsed.as_array().unwrap();
1309            assert!(arr[0].is_object(), "First element should be an object");
1310            assert!(
1311                arr[0]["_stream"].is_object(),
1312                "Should contain _stream object"
1313            );
1314            assert_eq!(arr[0]["_stream"]["origin"], "file:///test.txt");
1315            assert_eq!(
1316                arr[0]["_stream"]["placeholder"], true,
1317                "placeholder flag should be true"
1318            );
1319        }
1320    }
1321
1322    #[tokio::test]
1323    async fn test_aggregate_stream_bodies_with_none_origin() {
1324        use bytes::Bytes;
1325        use camel_api::{Body, StreamBody, StreamMetadata};
1326        use futures::stream;
1327        use tokio::sync::Mutex;
1328
1329        let chunks = vec![Ok(Bytes::from("test"))];
1330        let stream_body = StreamBody {
1331            stream: Arc::new(Mutex::new(Some(Box::pin(stream::iter(chunks))))),
1332            metadata: StreamMetadata {
1333                origin: None,
1334                ..Default::default()
1335            },
1336        };
1337
1338        let ex1 = Exchange::new(Message {
1339            headers: Default::default(),
1340            body: Body::Stream(stream_body),
1341        });
1342
1343        let exchanges = vec![ex1];
1344        let result = aggregate(exchanges, &AggregationStrategy::CollectAll);
1345
1346        let exchange = result.expect("Expected Ok result");
1347        assert!(
1348            matches!(exchange.input.body, Body::Json(_)),
1349            "Expected Json body"
1350        );
1351
1352        if let Body::Json(value) = exchange.input.body {
1353            let json_str = serde_json::to_string(&value).unwrap();
1354            let parsed: serde_json::Value = serde_json::from_str(&json_str).unwrap();
1355
1356            assert!(parsed.is_array(), "Result should be an array");
1357            let arr = parsed.as_array().unwrap();
1358            assert!(arr[0].is_object(), "First element should be an object");
1359            assert!(
1360                arr[0]["_stream"].is_object(),
1361                "Should contain _stream object"
1362            );
1363            assert_eq!(
1364                arr[0]["_stream"]["origin"],
1365                serde_json::Value::Null,
1366                "origin should be null when None"
1367            );
1368            assert_eq!(
1369                arr[0]["_stream"]["placeholder"], true,
1370                "placeholder flag should be true"
1371            );
1372        }
1373    }
1374
1375    #[tokio::test]
1376    async fn timeout_completion_clears_handle_from_map() {
1377        // Regression: Oracle audit found that natural timeout completion removed
1378        // the bucket but left the JoinHandle in `timeout_handles`, leaking the
1379        // entry until route shutdown. After fix, the timeout task itself cleans
1380        // its handle from the map on natural completion.
1381        let config = AggregatorConfig::correlate_by("key")
1382            .complete_on_timeout(Duration::from_millis(50))
1383            .build()
1384            .unwrap();
1385        let (tx, _rx) = mpsc::channel(256);
1386        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1387        let cancel = CancellationToken::new();
1388        let svc = AggregatorService::new(config, tx, registry, cancel);
1389
1390        // Send an exchange to create a pending bucket with a timeout task.
1391        let mut call_svc = svc.clone();
1392        let ex = make_exchange("key", "A", "data");
1393        let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1394        assert!(
1395            !svc.timeout_handles.lock().unwrap().is_empty(),
1396            "handle should exist while timeout pending"
1397        );
1398
1399        // Wait real time for the 50ms timeout to fire + spawned task to complete.
1400        tokio::time::sleep(Duration::from_millis(200)).await;
1401
1402        assert!(
1403            svc.timeout_handles.lock().unwrap().is_empty(),
1404            "handle should be cleared from map after natural timeout completion (was leak)"
1405        );
1406    }
1407
1408    #[tokio::test]
1409    async fn aggregator_shutdown_via_trait_dispatch() {
1410        // RED: builds an AggregatorService, dispatches through Arc<dyn StepLifecycle>,
1411        // and asserts idempotent shutdown works.
1412        let config = AggregatorConfig::correlate_by("key")
1413            .complete_when_size(10)
1414            .build()
1415            .unwrap();
1416        let (tx, _rx) = mpsc::channel(256);
1417        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1418        let cancel = CancellationToken::new();
1419        let svc = AggregatorService::new(config, tx, registry, cancel);
1420
1421        let step: Arc<dyn StepLifecycle> = Arc::new(svc);
1422        step.shutdown(StepShutdownReason::RouteStop)
1423            .await
1424            .expect("first shutdown should succeed");
1425        step.shutdown(StepShutdownReason::RouteStop)
1426            .await
1427            .expect("second shutdown (idempotent) should succeed");
1428    }
1429
1430    #[tokio::test(start_paused = true)]
1431    async fn test_shutdown_awaits_timeout_handles() {
1432        let config = AggregatorConfig::correlate_by("key")
1433            .complete_on_timeout(Duration::from_millis(100))
1434            .build()
1435            .unwrap();
1436        let (tx, _rx) = mpsc::channel(256);
1437        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1438        let cancel = CancellationToken::new();
1439        let svc = AggregatorService::new(config, tx, registry, cancel);
1440
1441        // Send an exchange to create a pending bucket with a timeout task.
1442        let mut call_svc = svc.clone();
1443        let ex = make_exchange("key", "A", "data");
1444        let _ = call_svc.ready().await.unwrap().call(ex).await.unwrap();
1445
1446        // Verify timeout handle exists.
1447        assert!(
1448            !svc.timeout_handles.lock().unwrap().is_empty(),
1449            "should have a timeout handle"
1450        );
1451
1452        // Shutdown should complete within the 5s deadline (the timeout task
1453        // gets cancelled so it won't wait for the full 100ms sleep).
1454        svc.shutdown_inner().await;
1455
1456        assert!(
1457            svc.timeout_handles.lock().unwrap().is_empty(),
1458            "all handles should be cleaned up after shutdown"
1459        );
1460    }
1461
1462    // ── R3-C1 Batch 1: DoS cap + background sweep ───────────────────
1463
1464    /// R3-C1: a flood of unique correlation keys must stay bounded.
1465    /// Default `max_buckets` is 10_000; the 10_001st unique key is rejected
1466    /// with `Aggregator reached maximum N buckets` (or its updated equivalent
1467    /// after the fix). The unique-key flood does NOT OOM the process.
1468    #[tokio::test]
1469    async fn test_unique_key_flood_stays_bounded_by_default() {
1470        // Builder defaults to max_buckets = 10_000, bucket_ttl = 300s.
1471        let config = AggregatorConfig::correlate_by("orderId")
1472            .complete_when_size(1_000_000) // never completes normally
1473            .build()
1474            .unwrap();
1475        let mut svc = new_test_svc(config);
1476
1477        // Send 10_001 unique keys. The first 10_000 should be accepted
1478        // (pending in their buckets); the 10_001st MUST be rejected.
1479        for i in 0..10_000usize {
1480            let ex = make_exchange("orderId", &format!("key-{i}"), "body");
1481            let result = svc.ready().await.unwrap().call(ex).await;
1482            assert!(result.is_ok(), "key {i} should be accepted under the cap");
1483        }
1484        let ex = make_exchange("orderId", "key-10001", "body");
1485        let result = svc.ready().await.unwrap().call(ex).await;
1486        assert!(
1487            result.is_err(),
1488            "10_001st unique key must be rejected by the max_buckets cap"
1489        );
1490        let err = result.unwrap_err().to_string();
1491        assert!(
1492            err.contains("maximum") || err.contains("max"),
1493            "error should mention cap: {err}"
1494        );
1495    }
1496
1497    /// The `AggregatorService` exposes a `sweep_handle` for the background
1498    /// sweep task. When `config.bucket_ttl` is `Some`, `AggregatorService::new`
1499    /// automatically spawns the sweep task and stores the handle, so the
1500    /// caller never sees `None` for a TTL-configured service. Cancelling the
1501    /// route token (via `shutdown`) aborts the sweep.
1502    #[tokio::test]
1503    async fn test_background_sweep_spawns_on_first_poll_not_construction() {
1504        let config = AggregatorConfig::correlate_by("key")
1505            .complete_when_size(10_000)
1506            .bucket_ttl(Duration::from_millis(50))
1507            .build()
1508            .unwrap();
1509        let cancel = CancellationToken::new();
1510        let (tx, _rx) = mpsc::channel(8);
1511        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1512        let mut svc = AggregatorService::new(config, tx, registry, cancel.clone());
1513
1514        // Construction is runtime-free: no sweep spawned yet.
1515        assert!(
1516            svc.sweep_handle
1517                .lock()
1518                .unwrap_or_else(|e| e.into_inner())
1519                .is_none(),
1520            "sweep must NOT be spawned at construction (runtime-free new)"
1521        );
1522
1523        // First readiness poll lazily spawns the sweep (a runtime is present
1524        // here). `ready()` drives `poll_ready` until Ready.
1525        let _ = svc.ready().await.unwrap();
1526        let sweep_present = svc
1527            .sweep_handle
1528            .lock()
1529            .unwrap_or_else(|e| e.into_inner())
1530            .is_some();
1531        assert!(
1532            sweep_present,
1533            "sweep handle should be Some after first poll when bucket_ttl is set"
1534        );
1535
1536        // Cancel the route token; the sweep task observes it and exits.
1537        cancel.cancel();
1538        // Give the task a moment to observe the cancel.
1539        tokio::time::sleep(Duration::from_millis(50)).await;
1540    }
1541
1542    // ── R3-M3: bounded timeout-task spawn ─────────────────────────────
1543
1544    #[tokio::test]
1545    async fn test_aggregator_timeout_task_cap_no_panic_under_flood() {
1546        // R3-M3: a flood of unique keys with a tiny max_timeout_tasks must not
1547        // spawn unbounded tasks, panic, or deadlock. Each call returns Ok(pending).
1548        use camel_api::aggregator::CorrelationStrategy;
1549
1550        let config = AggregatorConfig {
1551            header_name: "k".into(),
1552            completion: CompletionMode::Any(vec![
1553                CompletionCondition::Size(999),
1554                CompletionCondition::Timeout(Duration::from_secs(30)),
1555            ]),
1556            correlation: CorrelationStrategy::HeaderName("k".into()),
1557            strategy: AggregationStrategy::CollectAll,
1558            max_buckets: Some(50),
1559            bucket_ttl: Some(Duration::from_secs(30)),
1560            force_completion_on_stop: false,
1561            discard_on_timeout: false,
1562            max_timeout_tasks: 2,
1563        };
1564        let (late_tx, mut late_rx) = mpsc::channel(64);
1565        let registry: SharedLanguageRegistry = Arc::new(std::sync::Mutex::new(HashMap::new()));
1566        let cancel = CancellationToken::new();
1567        let svc = AggregatorService::new(config, late_tx, registry, cancel);
1568
1569        // Drive 20 unique-key exchanges — far exceeding max_timeout_tasks=2.
1570        for i in 0..20u64 {
1571            let mut ex = Exchange::new(Message {
1572                headers: HashMap::from([("k".to_string(), serde_json::json!(i))]),
1573                body: Body::Text(i.to_string()),
1574            });
1575            ex.input
1576                .headers
1577                .insert("k".to_string(), serde_json::json!(i));
1578            let outcome = tokio::time::timeout(Duration::from_secs(2), async {
1579                let mut s = svc.clone();
1580                use tower::ServiceExt;
1581                s.ready().await.unwrap().call(ex).await
1582            })
1583            .await;
1584            assert!(outcome.is_ok(), "call {} hung/panicked under task cap", i);
1585            // Each returns Ok(pending) since Size(999) is never reached.
1586            let res = outcome.unwrap().unwrap();
1587            assert_eq!(
1588                res.properties
1589                    .get(CAMEL_AGGREGATOR_PENDING)
1590                    .and_then(|v| v.as_bool()),
1591                Some(true),
1592                "exchange {} should be pending",
1593                i
1594            );
1595        }
1596        // Drain any late emissions to avoid blocking the channel.
1597        let _ = late_rx.try_recv();
1598    }
1599}