agent_framework_core/
middleware.rs

1//! Middleware pipelines for agents, chat clients, and function invocation.
2//!
3//! Rust equivalent of `agent_framework._middleware`. Middleware receives an
4//! owned context and a [`Next`] continuation. Call `next.run(ctx)` to continue
5//! the chain, mutate the context to observe/override results, or return the
6//! context directly (optionally with `terminate = true`) to short-circuit.
7
8use async_trait::async_trait;
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use crate::error::{Error, Result};
13use crate::tools::{BoxFuture, ToolDefinition};
14use crate::types::{AgentResponse, ChatOptions, ChatResponse, Message};
15
16/// The terminal handler invoked at the end of a middleware chain.
17pub type Terminal<C> = Box<dyn FnOnce(C) -> BoxFuture<Result<C>> + Send>;
18
19/// A middleware that transforms a context of type `C`.
20#[async_trait]
21pub trait Middleware<C: Send + 'static>: Send + Sync {
22    async fn process(&self, ctx: C, next: Next<C>) -> Result<C>;
23}
24
25/// The continuation passed to a [`Middleware`]. Calling [`Next::run`] invokes
26/// the remaining middleware and, finally, the terminal handler.
27pub struct Next<C: Send + 'static> {
28    middlewares: Arc<Vec<Arc<dyn Middleware<C>>>>,
29    index: usize,
30    terminal: Option<Terminal<C>>,
31}
32
33impl<C: Send + 'static> Next<C> {
34    /// Continue the chain with the given context.
35    pub async fn run(mut self, ctx: C) -> Result<C> {
36        if self.index < self.middlewares.len() {
37            let mw = self.middlewares[self.index].clone();
38            let next = Next {
39                middlewares: self.middlewares.clone(),
40                index: self.index + 1,
41                terminal: self.terminal.take(),
42            };
43            mw.process(ctx, next).await
44        } else if let Some(term) = self.terminal.take() {
45            term(ctx).await
46        } else {
47            Ok(ctx)
48        }
49    }
50}
51
52/// A pipeline of middleware of a single category.
53pub struct MiddlewarePipeline<C: Send + 'static> {
54    middlewares: Arc<Vec<Arc<dyn Middleware<C>>>>,
55}
56
57impl<C: Send + 'static> Default for MiddlewarePipeline<C> {
58    fn default() -> Self {
59        Self {
60            middlewares: Arc::new(Vec::new()),
61        }
62    }
63}
64
65impl<C: Send + 'static> Clone for MiddlewarePipeline<C> {
66    fn clone(&self) -> Self {
67        Self {
68            middlewares: self.middlewares.clone(),
69        }
70    }
71}
72
73impl<C: Send + 'static> MiddlewarePipeline<C> {
74    pub fn new(middlewares: Vec<Arc<dyn Middleware<C>>>) -> Self {
75        Self {
76            middlewares: Arc::new(middlewares),
77        }
78    }
79
80    pub fn is_empty(&self) -> bool {
81        self.middlewares.is_empty()
82    }
83
84    /// Execute the pipeline, running `terminal` after all middleware.
85    pub async fn execute(&self, ctx: C, terminal: Terminal<C>) -> Result<C> {
86        let next = Next {
87            middlewares: self.middlewares.clone(),
88            index: 0,
89            terminal: Some(terminal),
90        };
91        next.run(ctx).await
92    }
93}
94
95/// Context flowing through the agent middleware pipeline.
96pub struct AgentContext {
97    pub messages: Vec<Message>,
98    pub is_streaming: bool,
99    pub metadata: HashMap<String, serde_json::Value>,
100    /// The run result; populated by the terminal handler or overridden here.
101    pub result: Option<AgentResponse>,
102    /// If set to true, the pipeline stops without running further middleware.
103    pub terminate: bool,
104}
105
106impl AgentContext {
107    pub fn new(messages: Vec<Message>, is_streaming: bool) -> Self {
108        Self {
109            messages,
110            is_streaming,
111            metadata: HashMap::new(),
112            result: None,
113            terminate: false,
114        }
115    }
116}
117
118/// Context flowing through the chat middleware pipeline.
119pub struct ChatContext {
120    pub messages: Vec<Message>,
121    pub chat_options: ChatOptions,
122    pub is_streaming: bool,
123    pub metadata: HashMap<String, serde_json::Value>,
124    pub result: Option<ChatResponse>,
125    pub terminate: bool,
126}
127
128impl ChatContext {
129    pub fn new(messages: Vec<Message>, chat_options: ChatOptions, is_streaming: bool) -> Self {
130        Self {
131            messages,
132            chat_options,
133            is_streaming,
134            metadata: HashMap::new(),
135            result: None,
136            terminate: false,
137        }
138    }
139}
140
141/// The live, mutable tool list of an in-flight agent run (progressive tool
142/// exposure).
143///
144/// Handed to function middleware and tools via
145/// [`FunctionInvocationContext::tools`]; a clone is a view onto the *same*
146/// list. Mutations take effect on the **next** iteration of the
147/// function-calling loop — they never affect tool calls already requested in
148/// the in-flight batch, because the loop snapshots the list once per model
149/// iteration. Mirrors upstream `FunctionInvocationContext.tools` +
150/// `add_tools`/`remove_tools` (`_middleware.py`).
151#[derive(Clone, Default)]
152pub struct LiveToolList {
153    inner: Arc<std::sync::Mutex<Vec<ToolDefinition>>>,
154}
155
156impl LiveToolList {
157    /// A live list seeded with the run's current tools.
158    pub fn new(tools: Vec<ToolDefinition>) -> Self {
159        Self {
160            inner: Arc::new(std::sync::Mutex::new(tools)),
161        }
162    }
163
164    /// Add tools to the run (available to the model on the next iteration).
165    ///
166    /// Errors if any added tool's name collides with a tool already in the
167    /// list (mirrors upstream's `ValueError` on duplicate names); the whole
168    /// batch is validated first, so a duplicate leaves the list unchanged.
169    pub fn add_tools(&self, tools: impl IntoIterator<Item = ToolDefinition>) -> Result<()> {
170        let batch: Vec<ToolDefinition> = tools.into_iter().collect();
171        let mut list = self.inner.lock().unwrap();
172        for tool in &batch {
173            if list.iter().any(|t| t.name == tool.name)
174                || batch.iter().filter(|t| t.name == tool.name).count() > 1
175            {
176                return Err(Error::Configuration(format!(
177                    "cannot add tool '{}': a tool with that name already exists in this run",
178                    tool.name
179                )));
180            }
181        }
182        list.extend(batch);
183        Ok(())
184    }
185
186    /// Remove tools by name (effective on the next iteration). Names not
187    /// currently present are ignored.
188    pub fn remove_tools<'a>(&self, names: impl IntoIterator<Item = &'a str>) {
189        let to_remove: std::collections::HashSet<&str> = names.into_iter().collect();
190        self.inner
191            .lock()
192            .unwrap()
193            .retain(|t| !to_remove.contains(t.name.as_str()));
194    }
195
196    /// Whether a tool with `name` is currently in the list.
197    pub fn contains(&self, name: &str) -> bool {
198        self.inner.lock().unwrap().iter().any(|t| t.name == name)
199    }
200
201    /// A point-in-time copy of the list (what the next iteration will use).
202    pub fn snapshot(&self) -> Vec<ToolDefinition> {
203        self.inner.lock().unwrap().clone()
204    }
205}
206
207/// Context flowing through the function middleware pipeline.
208pub struct FunctionInvocationContext {
209    pub function_name: String,
210    pub arguments: serde_json::Value,
211    /// The [`AgentSession`](crate::session::AgentSession) of the agent run
212    /// this invocation belongs to, if the call originated from an agent run
213    /// with a session. Middleware may read it; tools receive it via
214    /// [`Tool::invoke_in_context`](crate::tools::Tool::invoke_in_context)
215    /// (the hook behind `Agent::as_tool` with `propagate_session`). Mirrors
216    /// upstream's `FunctionInvocationContext.session`.
217    pub session: Option<crate::session::AgentSession>,
218    /// The live, mutable tool list of the current agent run, or `None` when
219    /// the function is invoked outside a function-calling loop (e.g. via
220    /// [`Tool::invoke`](crate::tools::Tool::invoke) directly). Middleware and
221    /// tools may [`add_tools`](FunctionInvocationContext::add_tools) /
222    /// [`remove_tools`](FunctionInvocationContext::remove_tools); mutations
223    /// take effect on the **next** model iteration, not the in-flight batch.
224    pub tools: Option<LiveToolList>,
225    pub metadata: HashMap<String, serde_json::Value>,
226    pub result: Option<serde_json::Value>,
227    pub terminate: bool,
228}
229
230impl FunctionInvocationContext {
231    pub fn new(function_name: impl Into<String>, arguments: serde_json::Value) -> Self {
232        Self {
233            function_name: function_name.into(),
234            arguments,
235            session: None,
236            tools: None,
237            metadata: HashMap::new(),
238            result: None,
239            terminate: false,
240        }
241    }
242
243    /// Builder: attach the agent session this invocation belongs to.
244    pub fn with_session(mut self, session: Option<crate::session::AgentSession>) -> Self {
245        self.session = session;
246        self
247    }
248
249    /// Builder: attach the run's live tool list.
250    pub fn with_tools(mut self, tools: Option<LiveToolList>) -> Self {
251        self.tools = tools;
252        self
253    }
254
255    /// Add tools to the current agent run (progressive tool exposure); see
256    /// [`LiveToolList::add_tools`]. Errors when this invocation is not bound
257    /// to a live agent run (mirrors upstream's `RuntimeError`).
258    pub fn add_tools(&self, tools: impl IntoIterator<Item = ToolDefinition>) -> Result<()> {
259        self.tools
260            .as_ref()
261            .ok_or_else(|| {
262                Error::Configuration(
263                    "cannot add tools: this FunctionInvocationContext is not bound to a \
264                     live agent run"
265                        .into(),
266                )
267            })?
268            .add_tools(tools)
269    }
270
271    /// Remove tools from the current agent run by name; see
272    /// [`LiveToolList::remove_tools`]. Errors when this invocation is not
273    /// bound to a live agent run.
274    pub fn remove_tools<'a>(&self, names: impl IntoIterator<Item = &'a str>) -> Result<()> {
275        self.tools
276            .as_ref()
277            .ok_or_else(|| {
278                Error::Configuration(
279                    "cannot remove tools: this FunctionInvocationContext is not bound to a \
280                     live agent run"
281                        .into(),
282                )
283            })
284            .map(|t| t.remove_tools(names))
285    }
286}
287
288/// Convenience type aliases for each middleware category.
289pub type AgentMiddleware = dyn Middleware<AgentContext>;
290/// Chat middleware operates on a [`ChatContext`].
291pub type ChatMiddleware = dyn Middleware<ChatContext>;
292/// Function middleware operates on a [`FunctionInvocationContext`].
293pub type FunctionMiddleware = dyn Middleware<FunctionInvocationContext>;
294
295/// Adapter to build a [`Middleware`] from an async closure.
296pub struct FnMiddleware<C, F> {
297    f: F,
298    _marker: std::marker::PhantomData<fn(C)>,
299}
300
301impl<C, F> FnMiddleware<C, F> {
302    pub fn new(f: F) -> Self {
303        Self {
304            f,
305            _marker: std::marker::PhantomData,
306        }
307    }
308}
309
310#[async_trait]
311impl<C, F, Fut> Middleware<C> for FnMiddleware<C, F>
312where
313    C: Send + 'static,
314    F: Fn(C, Next<C>) -> Fut + Send + Sync,
315    Fut: std::future::Future<Output = Result<C>> + Send,
316{
317    async fn process(&self, ctx: C, next: Next<C>) -> Result<C> {
318        (self.f)(ctx, next).await
319    }
320}
agent_framework_core/middleware.rs

agent_framework_core/
middleware.rs
