ai_lib_rust/plugins/
middleware.rs

1//! Middleware system.
2
3use async_trait::async_trait;
4use std::sync::Arc;
5use crate::Result;
6
7#[derive(Debug, Clone)]
8pub struct MiddlewareContext {
9    pub request: serde_json::Value,
10    pub response: Option<serde_json::Value>,
11    pub request_id: Option<String>,
12    pub model: Option<String>,
13    pub metadata: std::collections::HashMap<String, serde_json::Value>,
14}
15
16impl MiddlewareContext {
17    pub fn new(request: serde_json::Value) -> Self { Self { request, response: None, request_id: None, model: None, metadata: std::collections::HashMap::new() } }
18    pub fn set_response(&mut self, r: serde_json::Value) { self.response = Some(r); }
19    pub fn with_request_id(mut self, id: impl Into<String>) -> Self { self.request_id = Some(id.into()); self }
20    pub fn with_model(mut self, m: impl Into<String>) -> Self { self.model = Some(m.into()); self }
21}
22
23pub type NextFn<'a> = Box<dyn FnOnce(MiddlewareContext) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<MiddlewareContext>> + Send + 'a>> + Send + 'a>;
24
25#[async_trait]
26pub trait Middleware: Send + Sync {
27    async fn process(&self, ctx: MiddlewareContext, next: NextFn<'_>) -> Result<MiddlewareContext>;
28    fn name(&self) -> &str { "unnamed" }
29}
30
31pub struct MiddlewareChain { middlewares: Vec<Arc<dyn Middleware>> }
32impl MiddlewareChain {
33    pub fn new() -> Self { Self { middlewares: Vec::new() } }
34    pub fn add(mut self, m: Arc<dyn Middleware>) -> Self { self.middlewares.push(m); self }
35    pub fn len(&self) -> usize { self.middlewares.len() }
36    pub fn is_empty(&self) -> bool { self.middlewares.is_empty() }
37
38    pub async fn execute<F, Fut>(&self, ctx: MiddlewareContext, handler: F) -> Result<MiddlewareContext>
39    where F: FnOnce(MiddlewareContext) -> Fut + Send + 'static, Fut: std::future::Future<Output = Result<MiddlewareContext>> + Send + 'static {
40        if self.middlewares.is_empty() { return handler(ctx).await; }
41        let mut current = ctx;
42        for mw in &self.middlewares {
43            let next: NextFn<'_> = Box::new(move |c| Box::pin(async move { Ok(c) }));
44            current = mw.process(current, next).await?;
45        }
46        handler(current).await
47    }
48}
49impl Default for MiddlewareChain { fn default() -> Self { Self::new() } }
ai_lib_rust/plugins/middleware.rs

ai_lib_rust/plugins/
middleware.rs
