gemini_adk_fluent_rs/compose/
middleware.rs

1//! M — Middleware composition.
2//!
3//! Stack middleware with `>>`: `M::a() >> M::b()` wraps `b` inside `a`, so
4//! order matters (the first layer sees a call first).
5//!
6//! ## Wiring status
7//!
8//! **TextAgent pipelines** (via `AgentBuilder::middleware` + `AgentBuilder::build`) —
9//! **fully wired**.  Every factory in this module produces a `MiddlewareComposite`
10//! whose layers are installed into the `LlmTextAgent` middleware chain at compile
11//! time.  Hooks fire in this order per `run()` call:
12//!
13//! 1. `before_model` (forward order) — may short-circuit with a cached response.
14//! 2. LLM call (skipped if `before_model` returned `Some`).
15//! 3. `after_model` (reverse order) — may replace the LLM response.
16//! 4. `before_tool` (forward) / `after_tool` (reverse) / `on_tool_error` (forward)
17//!    — called for each tool dispatch round.
18//! 5. `on_error` (forward) — called once if `run()` returns an error.
19//!
20//! **Live sessions** (via `Live::middleware`) — the **tool-lifecycle hooks**
21//! are wired: `before_tool` (a returned error vetoes the call), `after_tool`,
22//! and `on_tool_error` fire around every tool dispatch in the control lane,
23//! including background tools. Model-level hooks (`before_model`/`after_model`)
24//! do **not** apply to Live — a Live session streams over the wire and has no
25//! discrete `LlmRequest`/`LlmResponse` to intercept; use them on TextAgent
26//! pipelines instead.
27
28use std::sync::Arc;
29use std::time::Duration;
30
31use async_trait::async_trait;
32use gemini_genai_rs::prelude::FunctionCall;
33
34use gemini_adk_rs::context::AgentEvent;
35use gemini_adk_rs::error::{AgentError, ToolError};
36use gemini_adk_rs::middleware::{LatencyMiddleware, LogMiddleware, Middleware};
37
38/// A middleware composite — one or more middleware layers (`M::a() >> M::b()`).
39///
40/// A single `Arc<dyn Middleware>` converts into a one-layer composite, so
41/// `.middleware(Arc::new(MyLayer))` works without the namespace.
42#[derive(Clone)]
43#[non_exhaustive]
44pub struct MiddlewareComposite {
45    /// The ordered list of middleware layers.
46    pub layers: Vec<Arc<dyn Middleware>>,
47}
48
49impl MiddlewareComposite {
50    /// Create a composite containing a single middleware layer.
51    pub fn new(layer: Arc<dyn Middleware>) -> Self {
52        Self {
53            layers: vec![layer],
54        }
55    }
56
57    /// Number of layers.
58    pub fn len(&self) -> usize {
59        self.layers.len()
60    }
61
62    /// Whether empty.
63    pub fn is_empty(&self) -> bool {
64        self.layers.is_empty()
65    }
66}
67
68impl From<Arc<dyn Middleware>> for MiddlewareComposite {
69    fn from(layer: Arc<dyn Middleware>) -> Self {
70        Self::new(layer)
71    }
72}
73
74/// Stack two middleware composites with `>>`: `self` outside, `rhs` inside.
75impl std::ops::Shr for MiddlewareComposite {
76    type Output = MiddlewareComposite;
77
78    fn shr(mut self, rhs: MiddlewareComposite) -> Self::Output {
79        self.layers.extend(rhs.layers);
80        self
81    }
82}
83
84/// The `M` namespace — static factory methods for middleware.
85pub struct M;
86
87impl M {
88    /// Add logging middleware.
89    pub fn log() -> MiddlewareComposite {
90        MiddlewareComposite::new(Arc::new(LogMiddleware::new()))
91    }
92
93    /// Add latency tracking middleware.
94    pub fn latency() -> MiddlewareComposite {
95        MiddlewareComposite::new(Arc::new(LatencyMiddleware::new()))
96    }
97
98    /// Bound the agent run to `duration`. The text agent enforces the tightest
99    /// timeout across its middleware chain by wrapping the whole run; on elapse
100    /// it emits `AgentEvent::Timeout` and returns an error.
101    pub fn timeout(duration: Duration) -> MiddlewareComposite {
102        MiddlewareComposite::new(Arc::new(TimeoutMiddleware {
103            name: "timeout".to_string(),
104            duration,
105        }))
106    }
107
108    /// Add retry middleware — tracks errors and advises on retry.
109    pub fn retry(max_retries: u32) -> MiddlewareComposite {
110        MiddlewareComposite::new(Arc::new(gemini_adk_rs::middleware::RetryMiddleware::new(
111            max_retries,
112        )))
113    }
114
115    /// Add a custom event observer — called on every agent event.
116    pub fn tap(f: impl Fn(&AgentEvent) + Send + Sync + 'static) -> MiddlewareComposite {
117        MiddlewareComposite::new(Arc::new(TapMiddleware {
118            handler: Arc::new(f),
119        }))
120    }
121
122    /// Add a custom before-tool filter — called before every tool invocation.
123    pub fn before_tool(
124        f: impl Fn(&FunctionCall) -> Result<(), String> + Send + Sync + 'static,
125    ) -> MiddlewareComposite {
126        MiddlewareComposite::new(Arc::new(BeforeToolMiddleware {
127            handler: Arc::new(f),
128        }))
129    }
130
131    /// Add a custom after-tool hook — called after every successful tool invocation.
132    pub fn after_tool(
133        f: impl Fn(&FunctionCall, &serde_json::Value) -> Result<(), String> + Send + Sync + 'static,
134    ) -> MiddlewareComposite {
135        MiddlewareComposite::new(Arc::new(AfterToolMiddleware {
136            handler: Arc::new(f),
137        }))
138    }
139
140    /// Add a custom error observer — called when an agent-level error occurs.
141    pub fn on_error(
142        f: impl Fn(&AgentError) -> Result<(), String> + Send + Sync + 'static,
143    ) -> MiddlewareComposite {
144        MiddlewareComposite::new(Arc::new(OnErrorMiddleware {
145            handler: Arc::new(f),
146        }))
147    }
148
149    /// Add cost tracking middleware — records token usage estimates.
150    pub fn cost() -> MiddlewareComposite {
151        MiddlewareComposite::new(Arc::new(CostMiddleware {
152            tool_calls: std::sync::atomic::AtomicU64::new(0),
153        }))
154    }
155
156    /// Add rate-limiting middleware — spaces tool calls to at most `rps` per
157    /// second by delaying `before_tool` (concurrent calls queue rather than
158    /// burst).
159    pub fn rate_limit(rps: u32) -> MiddlewareComposite {
160        MiddlewareComposite::new(Arc::new(RateLimitMiddleware::new(rps)))
161    }
162
163    /// Add circuit breaker middleware — opens after consecutive failures.
164    pub fn circuit_breaker(threshold: u32) -> MiddlewareComposite {
165        MiddlewareComposite::new(Arc::new(CircuitBreakerMiddleware {
166            threshold,
167            consecutive_failures: std::sync::atomic::AtomicU32::new(0),
168        }))
169    }
170
171    /// Add tracing span middleware — creates spans for distributed tracing.
172    pub fn trace() -> MiddlewareComposite {
173        MiddlewareComposite::new(Arc::new(TraceMiddleware))
174    }
175
176    /// Add audit middleware — records all tool calls for review.
177    pub fn audit() -> MiddlewareComposite {
178        MiddlewareComposite::new(Arc::new(AuditMiddleware {
179            log: parking_lot::Mutex::new(Vec::new()),
180        }))
181    }
182
183    /// Scope middleware to specific agent names.
184    ///
185    /// Not yet enforced — agent-name routing requires dispatch-time filtering
186    /// the middleware chain doesn't expose, so this currently returns `inner`
187    /// unchanged. Hidden until real scoping lands to avoid implying behavior.
188    #[doc(hidden)]
189    pub fn scope(_names: &[&str], inner: MiddlewareComposite) -> MiddlewareComposite {
190        inner
191    }
192
193    /// Structured logging middleware — logs agent events as structured JSON.
194    pub fn structured_log() -> MiddlewareComposite {
195        MiddlewareComposite::new(Arc::new(StructuredLogMiddleware))
196    }
197
198    /// Dispatch logging middleware — logs dispatch/join events.
199    pub fn dispatch_log() -> MiddlewareComposite {
200        MiddlewareComposite::new(Arc::new(DispatchLogMiddleware))
201    }
202
203    /// Topology logging middleware — logs agent topology events.
204    pub fn topology_log() -> MiddlewareComposite {
205        MiddlewareComposite::new(Arc::new(TopologyLogMiddleware))
206    }
207
208    /// Add a tool input validator middleware.
209    pub fn validate(
210        f: impl Fn(&FunctionCall) -> Result<(), String> + Send + Sync + 'static,
211    ) -> MiddlewareComposite {
212        MiddlewareComposite::new(Arc::new(ValidateMiddleware {
213            validator: Arc::new(f),
214        }))
215    }
216
217    /// Fallback to an alternative model on error.
218    ///
219    /// Not yet enforced — swapping the model and retrying requires re-issuing
220    /// the LLM call, which the current `after_model`/`on_error` hooks can't do.
221    /// Hidden until real fallback lands to avoid implying behavior.
222    #[doc(hidden)]
223    pub fn fallback_model(model: &str) -> MiddlewareComposite {
224        MiddlewareComposite::new(Arc::new(FallbackModelMiddleware {
225            model: model.to_string(),
226        }))
227    }
228
229    /// Response caching middleware — caches model responses to avoid redundant calls.
230    pub fn cache() -> MiddlewareComposite {
231        MiddlewareComposite::new(Arc::new(CacheMiddleware {
232            cache: parking_lot::Mutex::new(std::collections::HashMap::new()),
233        }))
234    }
235
236    /// Deduplicate consecutive identical requests.
237    pub fn dedup() -> MiddlewareComposite {
238        MiddlewareComposite::new(Arc::new(DedupMiddleware {
239            last_request_hash: parking_lot::Mutex::new(None),
240        }))
241    }
242
243    /// Sample/pass-through a fraction of requests (0.0–1.0).
244    pub fn sample(rate: f64) -> MiddlewareComposite {
245        MiddlewareComposite::new(Arc::new(SampleMiddleware {
246            rate: rate.clamp(0.0, 1.0),
247        }))
248    }
249
250    /// Metrics collection middleware — tracks request counts, error counts, and latencies.
251    pub fn metrics() -> MiddlewareComposite {
252        MiddlewareComposite::new(Arc::new(MetricsMiddleware {
253            request_count: std::sync::atomic::AtomicU64::new(0),
254            error_count: std::sync::atomic::AtomicU64::new(0),
255        }))
256    }
257
258    /// Shortcut for a before-agent hook.
259    pub fn before_agent(
260        f: impl Fn(&gemini_adk_rs::context::InvocationContext) -> Result<(), String>
261        + Send
262        + Sync
263        + 'static,
264    ) -> MiddlewareComposite {
265        MiddlewareComposite::new(Arc::new(BeforeAgentMiddleware {
266            handler: Arc::new(f),
267        }))
268    }
269
270    /// Shortcut for an after-agent hook.
271    pub fn after_agent(
272        f: impl Fn(&gemini_adk_rs::context::InvocationContext) -> Result<(), String>
273        + Send
274        + Sync
275        + 'static,
276    ) -> MiddlewareComposite {
277        MiddlewareComposite::new(Arc::new(AfterAgentMiddleware {
278            handler: Arc::new(f),
279        }))
280    }
281
282    /// Shortcut for a before-model hook.
283    pub fn before_model(
284        f: impl Fn(&gemini_adk_rs::llm::LlmRequest) -> Result<(), String> + Send + Sync + 'static,
285    ) -> MiddlewareComposite {
286        MiddlewareComposite::new(Arc::new(BeforeModelMiddleware {
287            handler: Arc::new(f),
288        }))
289    }
290
291    /// Shortcut for an after-model hook.
292    pub fn after_model(
293        f: impl Fn(
294            &gemini_adk_rs::llm::LlmRequest,
295            &gemini_adk_rs::llm::LlmResponse,
296        ) -> Result<(), String>
297        + Send
298        + Sync
299        + 'static,
300    ) -> MiddlewareComposite {
301        MiddlewareComposite::new(Arc::new(AfterModelMiddleware {
302            handler: Arc::new(f),
303        }))
304    }
305
306    /// Loop iteration event hook — called on each iteration of a loop agent.
307    pub fn on_loop(f: impl Fn(u32) + Send + Sync + 'static) -> MiddlewareComposite {
308        MiddlewareComposite::new(Arc::new(OnLoopMiddleware {
309            handler: Arc::new(f),
310        }))
311    }
312
313    /// Timeout event hook — called when an agent times out.
314    pub fn on_timeout(f: impl Fn() + Send + Sync + 'static) -> MiddlewareComposite {
315        MiddlewareComposite::new(Arc::new(OnTimeoutMiddleware {
316            handler: Arc::new(f),
317        }))
318    }
319
320    /// Route decision event hook — called when a route agent selects a branch.
321    pub fn on_route(f: impl Fn(&str) + Send + Sync + 'static) -> MiddlewareComposite {
322        MiddlewareComposite::new(Arc::new(OnRouteMiddleware {
323            handler: Arc::new(f),
324        }))
325    }
326
327    /// Fallback event hook — called when a fallback agent activates.
328    pub fn on_fallback(f: impl Fn(&str) + Send + Sync + 'static) -> MiddlewareComposite {
329        MiddlewareComposite::new(Arc::new(OnFallbackMiddleware {
330            handler: Arc::new(f),
331        }))
332    }
333}
334
335/// Timeout middleware — stores the configured duration for runtime enforcement.
336#[allow(dead_code)]
337struct TimeoutMiddleware {
338    name: String,
339    duration: Duration,
340}
341
342#[async_trait::async_trait]
343impl Middleware for TimeoutMiddleware {
344    fn name(&self) -> &str {
345        &self.name
346    }
347
348    fn timeout(&self) -> Option<Duration> {
349        Some(self.duration)
350    }
351}
352
353// ── Tap Middleware ──────────────────────────────────────────────────────────
354
355struct TapMiddleware {
356    #[allow(clippy::type_complexity)]
357    handler: Arc<dyn Fn(&AgentEvent) + Send + Sync>,
358}
359
360#[async_trait]
361impl Middleware for TapMiddleware {
362    fn name(&self) -> &str {
363        "tap"
364    }
365
366    async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
367        (self.handler)(event);
368        Ok(())
369    }
370}
371
372// ── BeforeTool Middleware ───────────────────────────────────────────────────
373
374struct BeforeToolMiddleware {
375    #[allow(clippy::type_complexity)]
376    handler: Arc<dyn Fn(&FunctionCall) -> Result<(), String> + Send + Sync>,
377}
378
379#[async_trait]
380impl Middleware for BeforeToolMiddleware {
381    fn name(&self) -> &str {
382        "before_tool"
383    }
384
385    async fn before_tool(&self, call: &FunctionCall) -> Result<(), AgentError> {
386        (self.handler)(call).map_err(AgentError::Other)
387    }
388}
389
390// ── AfterTool Middleware ────────────────────────────────────────────────────
391
392struct AfterToolMiddleware {
393    #[allow(clippy::type_complexity)]
394    handler: Arc<dyn Fn(&FunctionCall, &serde_json::Value) -> Result<(), String> + Send + Sync>,
395}
396
397#[async_trait]
398impl Middleware for AfterToolMiddleware {
399    fn name(&self) -> &str {
400        "after_tool"
401    }
402
403    async fn after_tool(
404        &self,
405        call: &FunctionCall,
406        result: &serde_json::Value,
407    ) -> Result<(), AgentError> {
408        (self.handler)(call, result).map_err(AgentError::Other)
409    }
410}
411
412// ── OnError Middleware ──────────────────────────────────────────────────────
413
414struct OnErrorMiddleware {
415    #[allow(clippy::type_complexity)]
416    handler: Arc<dyn Fn(&AgentError) -> Result<(), String> + Send + Sync>,
417}
418
419#[async_trait]
420impl Middleware for OnErrorMiddleware {
421    fn name(&self) -> &str {
422        "on_error"
423    }
424
425    async fn on_error(&self, err: &AgentError) -> Result<(), AgentError> {
426        (self.handler)(err).map_err(AgentError::Other)
427    }
428}
429
430// ── Cost Middleware ────────────────────────────────────────────────────────
431
432/// Tracks the number of tool calls as a proxy for cost.
433pub struct CostMiddleware {
434    tool_calls: std::sync::atomic::AtomicU64,
435}
436
437impl CostMiddleware {
438    /// Returns the total number of tool calls recorded.
439    pub fn tool_call_count(&self) -> u64 {
440        self.tool_calls.load(std::sync::atomic::Ordering::SeqCst)
441    }
442}
443
444#[async_trait]
445impl Middleware for CostMiddleware {
446    fn name(&self) -> &str {
447        "cost"
448    }
449
450    async fn after_tool(
451        &self,
452        _call: &FunctionCall,
453        _result: &serde_json::Value,
454    ) -> Result<(), AgentError> {
455        self.tool_calls
456            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
457        Ok(())
458    }
459}
460
461// ── RateLimit Middleware ───────────────────────────────────────────────────
462
463#[allow(dead_code)]
464struct RateLimitMiddleware {
465    /// Minimum spacing between successive tool calls.
466    min_interval: Duration,
467    /// Reserved start time of the most recent call (advances as calls queue).
468    last: parking_lot::Mutex<Option<std::time::Instant>>,
469}
470
471impl RateLimitMiddleware {
472    fn new(rps: u32) -> Self {
473        let rps = rps.max(1);
474        Self {
475            min_interval: Duration::from_secs_f64(1.0 / rps as f64),
476            last: parking_lot::Mutex::new(None),
477        }
478    }
479}
480
481#[async_trait]
482impl Middleware for RateLimitMiddleware {
483    fn name(&self) -> &str {
484        "rate_limit"
485    }
486
487    async fn before_tool(&self, _call: &FunctionCall) -> Result<(), AgentError> {
488        // Compute the wait needed to honor min_interval, reserving this call's
489        // slot so concurrent callers queue instead of bursting. The lock is
490        // released before the await (never held across it).
491        let wait = {
492            let mut last = self.last.lock();
493            let now = std::time::Instant::now();
494            let scheduled = match *last {
495                Some(prev) if prev + self.min_interval > now => prev + self.min_interval,
496                _ => now,
497            };
498            *last = Some(scheduled);
499            scheduled.saturating_duration_since(now)
500        };
501        if !wait.is_zero() {
502            tokio::time::sleep(wait).await;
503        }
504        Ok(())
505    }
506}
507
508// ── CircuitBreaker Middleware ──────────────────────────────────────────────
509
510struct CircuitBreakerMiddleware {
511    threshold: u32,
512    consecutive_failures: std::sync::atomic::AtomicU32,
513}
514
515#[async_trait]
516impl Middleware for CircuitBreakerMiddleware {
517    fn name(&self) -> &str {
518        "circuit_breaker"
519    }
520
521    async fn before_tool(&self, _call: &FunctionCall) -> Result<(), AgentError> {
522        let failures = self
523            .consecutive_failures
524            .load(std::sync::atomic::Ordering::SeqCst);
525        if failures >= self.threshold {
526            return Err(AgentError::Other(format!(
527                "Circuit breaker open: {} consecutive failures (threshold: {})",
528                failures, self.threshold
529            )));
530        }
531        Ok(())
532    }
533
534    async fn after_tool(
535        &self,
536        _call: &FunctionCall,
537        _result: &serde_json::Value,
538    ) -> Result<(), AgentError> {
539        self.consecutive_failures
540            .store(0, std::sync::atomic::Ordering::SeqCst);
541        Ok(())
542    }
543
544    async fn on_tool_error(
545        &self,
546        _call: &FunctionCall,
547        _err: &ToolError,
548    ) -> Result<(), AgentError> {
549        self.consecutive_failures
550            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
551        Ok(())
552    }
553}
554
555// ── Trace Middleware ──────────────────────────────────────────────────────
556
557/// Middleware that creates tracing spans for agent and tool lifecycle events.
558/// With an OTel exporter feature enabled, these spans are picked up by
559/// `tracing-opentelemetry` and exported as OTel spans.
560struct TraceMiddleware;
561
562#[async_trait]
563impl Middleware for TraceMiddleware {
564    fn name(&self) -> &str {
565        "trace"
566    }
567
568    async fn before_agent(
569        &self,
570        ctx: &gemini_adk_rs::context::InvocationContext,
571    ) -> Result<(), AgentError> {
572        let sid = ctx.session_id.as_deref().unwrap_or("unknown");
573        gemini_adk_rs::telemetry::logging::log_agent_started(sid, 0);
574        Ok(())
575    }
576
577    async fn before_tool(&self, call: &FunctionCall) -> Result<(), AgentError> {
578        gemini_adk_rs::telemetry::logging::log_tool_dispatch("fluent", &call.name, "function");
579        Ok(())
580    }
581
582    async fn after_tool(
583        &self,
584        call: &FunctionCall,
585        _result: &serde_json::Value,
586    ) -> Result<(), AgentError> {
587        gemini_adk_rs::telemetry::logging::log_tool_result("fluent", &call.name, true, 0.0);
588        Ok(())
589    }
590
591    async fn on_tool_error(&self, call: &FunctionCall, _err: &ToolError) -> Result<(), AgentError> {
592        gemini_adk_rs::telemetry::logging::log_tool_result("fluent", &call.name, false, 0.0);
593        Ok(())
594    }
595
596    async fn on_error(&self, err: &AgentError) -> Result<(), AgentError> {
597        gemini_adk_rs::telemetry::logging::log_agent_error("fluent", &err.to_string());
598        Ok(())
599    }
600}
601
602// ── Audit Middleware ─────────────────────────────────────────────────────
603
604/// Records all tool calls for audit review.
605pub struct AuditMiddleware {
606    log: parking_lot::Mutex<Vec<AuditEntry>>,
607}
608
609/// An audit log entry.
610#[derive(Debug, Clone)]
611pub struct AuditEntry {
612    /// Tool name.
613    pub tool_name: String,
614    /// Tool arguments.
615    pub args: serde_json::Value,
616    /// Whether the call succeeded.
617    pub success: Option<bool>,
618}
619
620impl AuditMiddleware {
621    /// Returns a snapshot of the audit log.
622    pub fn entries(&self) -> Vec<AuditEntry> {
623        self.log.lock().clone()
624    }
625}
626
627#[async_trait]
628impl Middleware for AuditMiddleware {
629    fn name(&self) -> &str {
630        "audit"
631    }
632
633    async fn before_tool(&self, call: &FunctionCall) -> Result<(), AgentError> {
634        let mut log = self.log.lock();
635        if log.len() >= 10_000 {
636            log.drain(..1_000);
637        }
638        log.push(AuditEntry {
639            tool_name: call.name.clone(),
640            args: call.args.clone(),
641            success: None,
642        });
643        Ok(())
644    }
645
646    async fn after_tool(
647        &self,
648        call: &FunctionCall,
649        _result: &serde_json::Value,
650    ) -> Result<(), AgentError> {
651        let mut log = self.log.lock();
652        if let Some(entry) = log.iter_mut().rev().find(|e| e.tool_name == call.name) {
653            entry.success = Some(true);
654        }
655        Ok(())
656    }
657
658    async fn on_tool_error(&self, call: &FunctionCall, _err: &ToolError) -> Result<(), AgentError> {
659        let mut log = self.log.lock();
660        if let Some(entry) = log.iter_mut().rev().find(|e| e.tool_name == call.name) {
661            entry.success = Some(false);
662        }
663        Ok(())
664    }
665}
666
667// ── Validate Middleware ──────────────────────────────────────────────────
668
669struct ValidateMiddleware {
670    #[allow(clippy::type_complexity)]
671    validator: Arc<dyn Fn(&FunctionCall) -> Result<(), String> + Send + Sync>,
672}
673
674#[async_trait]
675impl Middleware for ValidateMiddleware {
676    fn name(&self) -> &str {
677        "validate"
678    }
679
680    async fn before_tool(&self, call: &FunctionCall) -> Result<(), AgentError> {
681        (self.validator)(call).map_err(|e| AgentError::Tool(ToolError::InvalidArgs(e)))
682    }
683}
684
685// ── FallbackModel Middleware ──────────────────────────────────────────
686
687/// Middleware that falls back to an alternative model on error.
688#[allow(dead_code)]
689struct FallbackModelMiddleware {
690    model: String,
691}
692
693#[async_trait]
694impl Middleware for FallbackModelMiddleware {
695    fn name(&self) -> &str {
696        "fallback_model"
697    }
698
699    async fn on_error(&self, _err: &AgentError) -> Result<(), AgentError> {
700        // Runtime inspects the `model` field and retries with the fallback model.
701        Ok(())
702    }
703}
704
705// ── Cache Middleware ──────────────────────────────────────────────────
706
707/// Caches model responses keyed by request hash to avoid redundant LLM calls.
708pub struct CacheMiddleware {
709    cache: parking_lot::Mutex<std::collections::HashMap<u64, gemini_adk_rs::llm::LlmResponse>>,
710}
711
712impl CacheMiddleware {
713    /// Returns the number of cached entries.
714    pub fn len(&self) -> usize {
715        self.cache.lock().len()
716    }
717
718    /// Whether the cache is empty.
719    pub fn is_empty(&self) -> bool {
720        self.cache.lock().is_empty()
721    }
722
723    /// Clear all cached entries.
724    pub fn clear(&self) {
725        self.cache.lock().clear();
726    }
727}
728
729#[async_trait]
730impl Middleware for CacheMiddleware {
731    fn name(&self) -> &str {
732        "cache"
733    }
734
735    async fn before_model(
736        &self,
737        request: &gemini_adk_rs::llm::LlmRequest,
738    ) -> Result<Option<gemini_adk_rs::llm::LlmResponse>, AgentError> {
739        use std::hash::{Hash, Hasher};
740        let mut hasher = std::collections::hash_map::DefaultHasher::new();
741        format!("{request:?}").hash(&mut hasher);
742        let key = hasher.finish();
743        let cache = self.cache.lock();
744        Ok(cache.get(&key).cloned())
745    }
746
747    async fn after_model(
748        &self,
749        request: &gemini_adk_rs::llm::LlmRequest,
750        response: &gemini_adk_rs::llm::LlmResponse,
751    ) -> Result<Option<gemini_adk_rs::llm::LlmResponse>, AgentError> {
752        use std::hash::{Hash, Hasher};
753        let mut hasher = std::collections::hash_map::DefaultHasher::new();
754        format!("{request:?}").hash(&mut hasher);
755        let key = hasher.finish();
756        self.cache.lock().insert(key, response.clone());
757        Ok(None) // don't replace the response
758    }
759}
760
761// ── Dedup Middleware ─────────────────────────────────────────────────
762
763/// Deduplicates consecutive identical requests by hashing.
764#[allow(dead_code)]
765struct DedupMiddleware {
766    last_request_hash: parking_lot::Mutex<Option<u64>>,
767}
768
769#[async_trait]
770impl Middleware for DedupMiddleware {
771    fn name(&self) -> &str {
772        "dedup"
773    }
774
775    async fn before_model(
776        &self,
777        request: &gemini_adk_rs::llm::LlmRequest,
778    ) -> Result<Option<gemini_adk_rs::llm::LlmResponse>, AgentError> {
779        use std::hash::{Hash, Hasher};
780        let mut hasher = std::collections::hash_map::DefaultHasher::new();
781        format!("{request:?}").hash(&mut hasher);
782        let hash = hasher.finish();
783        let mut last = self.last_request_hash.lock();
784        if *last == Some(hash) {
785            // Duplicate consecutive request — signal skip by returning empty response.
786            return Err(AgentError::Other(
787                "Duplicate consecutive request".to_string(),
788            ));
789        }
790        *last = Some(hash);
791        Ok(None)
792    }
793}
794
795// ── Sample Middleware ────────────────────────────────────────────────
796
797/// Passes through only a fraction of requests, dropping the rest.
798#[allow(dead_code)]
799struct SampleMiddleware {
800    rate: f64,
801}
802
803#[async_trait]
804impl Middleware for SampleMiddleware {
805    fn name(&self) -> &str {
806        "sample"
807    }
808
809    async fn before_model(
810        &self,
811        _request: &gemini_adk_rs::llm::LlmRequest,
812    ) -> Result<Option<gemini_adk_rs::llm::LlmResponse>, AgentError> {
813        use std::hash::{Hash, Hasher};
814        // Use a fast pseudo-random check based on time.
815        let mut hasher = std::collections::hash_map::DefaultHasher::new();
816        std::time::Instant::now().hash(&mut hasher);
817        let hash = hasher.finish();
818        let normalized = (hash as f64) / (u64::MAX as f64);
819        if normalized > self.rate {
820            return Err(AgentError::Other("Sampled out".to_string()));
821        }
822        Ok(None)
823    }
824}
825
826// ── Metrics Middleware ──────────────────────────────────────────────
827
828/// Collects request and error counts.
829pub struct MetricsMiddleware {
830    request_count: std::sync::atomic::AtomicU64,
831    error_count: std::sync::atomic::AtomicU64,
832}
833
834impl MetricsMiddleware {
835    /// Returns the total number of requests observed.
836    pub fn request_count(&self) -> u64 {
837        self.request_count.load(std::sync::atomic::Ordering::SeqCst)
838    }
839
840    /// Returns the total number of errors observed.
841    pub fn error_count(&self) -> u64 {
842        self.error_count.load(std::sync::atomic::Ordering::SeqCst)
843    }
844}
845
846#[async_trait]
847impl Middleware for MetricsMiddleware {
848    fn name(&self) -> &str {
849        "metrics"
850    }
851
852    async fn before_agent(
853        &self,
854        _ctx: &gemini_adk_rs::context::InvocationContext,
855    ) -> Result<(), AgentError> {
856        self.request_count
857            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
858        Ok(())
859    }
860
861    async fn on_error(&self, _err: &AgentError) -> Result<(), AgentError> {
862        self.error_count
863            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
864        Ok(())
865    }
866}
867
868// ── BeforeAgent Middleware ───────────────────────────────────────────
869
870struct BeforeAgentMiddleware {
871    #[allow(clippy::type_complexity)]
872    handler:
873        Arc<dyn Fn(&gemini_adk_rs::context::InvocationContext) -> Result<(), String> + Send + Sync>,
874}
875
876#[async_trait]
877impl Middleware for BeforeAgentMiddleware {
878    fn name(&self) -> &str {
879        "before_agent"
880    }
881
882    async fn before_agent(
883        &self,
884        ctx: &gemini_adk_rs::context::InvocationContext,
885    ) -> Result<(), AgentError> {
886        (self.handler)(ctx).map_err(AgentError::Other)
887    }
888}
889
890// ── AfterAgent Middleware ───────────────────────────────────────────
891
892struct AfterAgentMiddleware {
893    #[allow(clippy::type_complexity)]
894    handler:
895        Arc<dyn Fn(&gemini_adk_rs::context::InvocationContext) -> Result<(), String> + Send + Sync>,
896}
897
898#[async_trait]
899impl Middleware for AfterAgentMiddleware {
900    fn name(&self) -> &str {
901        "after_agent"
902    }
903
904    async fn after_agent(
905        &self,
906        ctx: &gemini_adk_rs::context::InvocationContext,
907    ) -> Result<(), AgentError> {
908        (self.handler)(ctx).map_err(AgentError::Other)
909    }
910}
911
912// ── BeforeModel Middleware ──────────────────────────────────────────
913
914struct BeforeModelMiddleware {
915    #[allow(clippy::type_complexity)]
916    handler: Arc<dyn Fn(&gemini_adk_rs::llm::LlmRequest) -> Result<(), String> + Send + Sync>,
917}
918
919#[async_trait]
920impl Middleware for BeforeModelMiddleware {
921    fn name(&self) -> &str {
922        "before_model"
923    }
924
925    async fn before_model(
926        &self,
927        request: &gemini_adk_rs::llm::LlmRequest,
928    ) -> Result<Option<gemini_adk_rs::llm::LlmResponse>, AgentError> {
929        (self.handler)(request).map_err(AgentError::Other)?;
930        Ok(None)
931    }
932}
933
934// ── AfterModel Middleware ──────────────────────────────────────────
935
936struct AfterModelMiddleware {
937    #[allow(clippy::type_complexity)]
938    handler: Arc<
939        dyn Fn(
940                &gemini_adk_rs::llm::LlmRequest,
941                &gemini_adk_rs::llm::LlmResponse,
942            ) -> Result<(), String>
943            + Send
944            + Sync,
945    >,
946}
947
948#[async_trait]
949impl Middleware for AfterModelMiddleware {
950    fn name(&self) -> &str {
951        "after_model"
952    }
953
954    async fn after_model(
955        &self,
956        request: &gemini_adk_rs::llm::LlmRequest,
957        response: &gemini_adk_rs::llm::LlmResponse,
958    ) -> Result<Option<gemini_adk_rs::llm::LlmResponse>, AgentError> {
959        (self.handler)(request, response).map_err(AgentError::Other)?;
960        Ok(None)
961    }
962}
963
964// ── OnLoop Middleware ───────────────────────────────────────────────
965
966struct OnLoopMiddleware {
967    handler: Arc<dyn Fn(u32) + Send + Sync>,
968}
969
970#[async_trait]
971impl Middleware for OnLoopMiddleware {
972    fn name(&self) -> &str {
973        "on_loop"
974    }
975
976    async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
977        if let AgentEvent::LoopIteration { iteration } = event {
978            (self.handler)(*iteration);
979        }
980        Ok(())
981    }
982}
983
984// ── OnTimeout Middleware ────────────────────────────────────────────
985
986struct OnTimeoutMiddleware {
987    handler: Arc<dyn Fn() + Send + Sync>,
988}
989
990#[async_trait]
991impl Middleware for OnTimeoutMiddleware {
992    fn name(&self) -> &str {
993        "on_timeout"
994    }
995
996    async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
997        if let AgentEvent::Timeout = event {
998            (self.handler)();
999        }
1000        Ok(())
1001    }
1002}
1003
1004// ── OnRoute Middleware ──────────────────────────────────────────────
1005
1006struct OnRouteMiddleware {
1007    handler: Arc<dyn Fn(&str) + Send + Sync>,
1008}
1009
1010#[async_trait]
1011impl Middleware for OnRouteMiddleware {
1012    fn name(&self) -> &str {
1013        "on_route"
1014    }
1015
1016    async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
1017        if let AgentEvent::RouteSelected { agent_name } = event {
1018            (self.handler)(agent_name);
1019        }
1020        Ok(())
1021    }
1022}
1023
1024// ── OnFallback Middleware ───────────────────────────────────────────
1025
1026struct OnFallbackMiddleware {
1027    handler: Arc<dyn Fn(&str) + Send + Sync>,
1028}
1029
1030#[async_trait]
1031impl Middleware for OnFallbackMiddleware {
1032    fn name(&self) -> &str {
1033        "on_fallback"
1034    }
1035
1036    async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
1037        if let AgentEvent::FallbackActivated { agent_name } = event {
1038            (self.handler)(agent_name);
1039        }
1040        Ok(())
1041    }
1042}
1043
1044// ── Structured Log Middleware ────────────────────────────────────────
1045
1046struct StructuredLogMiddleware;
1047
1048#[async_trait]
1049impl Middleware for StructuredLogMiddleware {
1050    fn name(&self) -> &str {
1051        "structured_log"
1052    }
1053
1054    async fn on_event(&self, event: &AgentEvent) -> Result<(), AgentError> {
1055        // Log events as structured format (uses tracing in production).
1056        let _ = event;
1057        Ok(())
1058    }
1059}
1060
1061// ── Dispatch Log Middleware ──────────────────────────────────────────
1062
1063struct DispatchLogMiddleware;
1064
1065#[async_trait]
1066impl Middleware for DispatchLogMiddleware {
1067    fn name(&self) -> &str {
1068        "dispatch_log"
1069    }
1070}
1071
1072// ── Topology Log Middleware ──────────────────────────────────────────
1073
1074struct TopologyLogMiddleware;
1075
1076#[async_trait]
1077impl Middleware for TopologyLogMiddleware {
1078    fn name(&self) -> &str {
1079        "topology_log"
1080    }
1081}
1082
1083#[cfg(test)]
1084mod tests {
1085    use super::*;
1086
1087    #[test]
1088    fn log_creates_composite() {
1089        let m = M::log();
1090        assert_eq!(m.len(), 1);
1091    }
1092
1093    #[test]
1094    fn latency_creates_composite() {
1095        let m = M::latency();
1096        assert_eq!(m.len(), 1);
1097    }
1098
1099    #[test]
1100    fn timeout_creates_composite() {
1101        let m = M::timeout(Duration::from_secs(30));
1102        assert_eq!(m.len(), 1);
1103    }
1104
1105    #[test]
1106    fn compose_with_bitor() {
1107        let m = M::log() >> M::latency() >> M::timeout(Duration::from_secs(5));
1108        assert_eq!(m.len(), 3);
1109    }
1110
1111    #[test]
1112    fn retry_creates_composite() {
1113        let m = M::retry(3);
1114        assert_eq!(m.len(), 1);
1115    }
1116
1117    #[test]
1118    fn tap_creates_composite() {
1119        let m = M::tap(|_event| {});
1120        assert_eq!(m.len(), 1);
1121    }
1122
1123    #[test]
1124    fn before_tool_creates_composite() {
1125        let m = M::before_tool(|_call| Ok(()));
1126        assert_eq!(m.len(), 1);
1127    }
1128
1129    #[test]
1130    fn cost_creates_composite() {
1131        let m = M::cost();
1132        assert_eq!(m.len(), 1);
1133    }
1134
1135    #[test]
1136    fn rate_limit_creates_composite() {
1137        let m = M::rate_limit(10);
1138        assert_eq!(m.len(), 1);
1139    }
1140
1141    #[test]
1142    fn circuit_breaker_creates_composite() {
1143        let m = M::circuit_breaker(5);
1144        assert_eq!(m.len(), 1);
1145    }
1146
1147    #[test]
1148    fn trace_creates_composite() {
1149        let m = M::trace();
1150        assert_eq!(m.len(), 1);
1151    }
1152
1153    #[test]
1154    fn audit_creates_composite() {
1155        let m = M::audit();
1156        assert_eq!(m.len(), 1);
1157    }
1158
1159    #[test]
1160    fn validate_creates_composite() {
1161        let m = M::validate(|_call| Ok(()));
1162        assert_eq!(m.len(), 1);
1163    }
1164
1165    #[test]
1166    fn fallback_model_creates_composite() {
1167        let m = M::fallback_model("gemini-1.5-flash");
1168        assert_eq!(m.len(), 1);
1169    }
1170
1171    #[test]
1172    fn cache_creates_composite() {
1173        let m = M::cache();
1174        assert_eq!(m.len(), 1);
1175    }
1176
1177    #[test]
1178    fn dedup_creates_composite() {
1179        let m = M::dedup();
1180        assert_eq!(m.len(), 1);
1181    }
1182
1183    #[test]
1184    fn sample_creates_composite() {
1185        let m = M::sample(0.5);
1186        assert_eq!(m.len(), 1);
1187    }
1188
1189    #[test]
1190    fn sample_clamps_rate() {
1191        let m = M::sample(2.0);
1192        assert_eq!(m.len(), 1);
1193        let m = M::sample(-1.0);
1194        assert_eq!(m.len(), 1);
1195    }
1196
1197    #[test]
1198    fn metrics_creates_composite() {
1199        let m = M::metrics();
1200        assert_eq!(m.len(), 1);
1201    }
1202
1203    #[test]
1204    fn before_agent_creates_composite() {
1205        let m = M::before_agent(|_ctx| Ok(()));
1206        assert_eq!(m.len(), 1);
1207    }
1208
1209    #[test]
1210    fn after_agent_creates_composite() {
1211        let m = M::after_agent(|_ctx| Ok(()));
1212        assert_eq!(m.len(), 1);
1213    }
1214
1215    #[test]
1216    fn before_model_creates_composite() {
1217        let m = M::before_model(|_req| Ok(()));
1218        assert_eq!(m.len(), 1);
1219    }
1220
1221    #[test]
1222    fn after_model_creates_composite() {
1223        let m = M::after_model(|_req, _resp| Ok(()));
1224        assert_eq!(m.len(), 1);
1225    }
1226
1227    #[test]
1228    fn on_loop_creates_composite() {
1229        let m = M::on_loop(|_iteration| {});
1230        assert_eq!(m.len(), 1);
1231    }
1232
1233    #[test]
1234    fn on_timeout_creates_composite() {
1235        let m = M::on_timeout(|| {});
1236        assert_eq!(m.len(), 1);
1237    }
1238
1239    #[test]
1240    fn on_route_creates_composite() {
1241        let m = M::on_route(|_name| {});
1242        assert_eq!(m.len(), 1);
1243    }
1244
1245    #[test]
1246    fn on_fallback_creates_composite() {
1247        let m = M::on_fallback(|_name| {});
1248        assert_eq!(m.len(), 1);
1249    }
1250
1251    #[test]
1252    fn compose_all_middleware() {
1253        let m = M::log()
1254            >> M::latency()
1255            >> M::timeout(Duration::from_secs(30))
1256            >> M::retry(3)
1257            >> M::cost()
1258            >> M::rate_limit(10)
1259            >> M::circuit_breaker(5)
1260            >> M::trace()
1261            >> M::audit()
1262            >> M::validate(|_| Ok(()))
1263            >> M::fallback_model("gemini-1.5-flash")
1264            >> M::cache()
1265            >> M::dedup()
1266            >> M::sample(0.5)
1267            >> M::metrics()
1268            >> M::before_agent(|_| Ok(()))
1269            >> M::after_agent(|_| Ok(()))
1270            >> M::before_model(|_| Ok(()))
1271            >> M::after_model(|_, _| Ok(()))
1272            >> M::on_loop(|_| {})
1273            >> M::on_timeout(|| {})
1274            >> M::on_route(|_| {})
1275            >> M::on_fallback(|_| {});
1276        assert_eq!(m.len(), 23);
1277    }
1278}
gemini_adk_fluent_rs/compose/middleware.rs

gemini_adk_fluent_rs/compose/
middleware.rs
