agentlib_core/
middleware.rs

1use crate::types::{ExecutionContext, MiddlewareScope};
2use anyhow::Result;
3use async_trait::async_trait;
4
5pub struct MiddlewareContext<'a> {
6    pub scope: MiddlewareScope,
7    pub ctx: &'a mut ExecutionContext,
8    pub tool: Option<ToolMiddlewareContext>,
9}
10
11#[derive(Debug, Clone)]
12pub struct ToolMiddlewareContext {
13    pub name: String,
14    pub args: serde_json::Value,
15    pub result: Option<serde_json::Value>,
16}
17
18#[async_trait]
19pub trait Middleware: Send + Sync {
20    fn name(&self) -> &str;
21    async fn run(
22        &self,
23        m_ctx: &mut MiddlewareContext<'_>,
24        next: &mut dyn NextMiddleware,
25    ) -> Result<()>;
26}
27
28#[async_trait]
29pub trait NextMiddleware: Send {
30    async fn run(&mut self, m_ctx: &mut MiddlewareContext<'_>) -> Result<()>;
31}
32
33pub struct MiddlewarePipeline {
34    middlewares: Vec<Box<dyn Middleware>>,
35}
36
37impl MiddlewarePipeline {
38    pub fn new() -> Self {
39        Self {
40            middlewares: Vec::new(),
41        }
42    }
43
44    pub fn use_middleware(&mut self, middleware: Box<dyn Middleware>) {
45        self.middlewares.push(middleware);
46    }
47
48    pub async fn run(&self, scope: MiddlewareScope, ctx: &mut ExecutionContext) -> Result<()> {
49        let mut m_ctx = MiddlewareContext {
50            scope,
51            ctx,
52            tool: None,
53        };
54        let mut runner = PipelineRunner {
55            middlewares: &self.middlewares,
56            index: 0,
57        };
58        runner.run(&mut m_ctx).await
59    }
60
61    pub async fn run_tool(
62        &self,
63        scope: MiddlewareScope,
64        ctx: &mut ExecutionContext,
65        tool: ToolMiddlewareContext,
66    ) -> Result<()> {
67        let mut m_ctx = MiddlewareContext {
68            scope,
69            ctx,
70            tool: Some(tool),
71        };
72        let mut runner = PipelineRunner {
73            middlewares: &self.middlewares,
74            index: 0,
75        };
76        runner.run(&mut m_ctx).await
77    }
78}
79
80struct PipelineRunner<'a> {
81    middlewares: &'a [Box<dyn Middleware>],
82    index: usize,
83}
84
85#[async_trait]
86impl<'a> NextMiddleware for PipelineRunner<'a> {
87    async fn run(&mut self, m_ctx: &mut MiddlewareContext<'_>) -> Result<()> {
88        if self.index < self.middlewares.len() {
89            let middleware = &self.middlewares[self.index];
90            self.index += 1;
91            middleware.run(m_ctx, self).await
92        } else {
93            Ok(())
94        }
95    }
96}
agentlib_core/middleware.rs

agentlib_core/
middleware.rs
