ag_ui/axum/
router.rs

1//! Mounting an agent on a router.
2//!
3//! ```
4//! use ag_ui::axum::RouterExt;
5//! use axum::Router;
6//! use axum::routing::get;
7//! # use ag_ui::RunOutcome;
8//! # use ag_ui::server::{Agent, Result, RunContext};
9//! # struct CartAgent;
10//! # impl Agent for CartAgent {
11//! #     type State = ();
12//! #     async fn run(&self, _ctx: &mut RunContext<()>) -> Result<RunOutcome> { Ok(RunOutcome::Success) }
13//! # }
14//!
15//! let app: Router = Router::new()
16//!     .route("/health", get(|| async { "ok" }))
17//!     .route_agui("/agent", CartAgent);
18//! # let _ = app;
19//! ```
20
21use std::sync::Arc;
22use std::time::Duration;
23
24use crate::server::{Agent, Runner, StreamTransformer, TransformerChain};
25use axum::Router;
26use axum::http::{HeaderMap, header};
27use axum::response::{IntoResponse, Response};
28use axum::routing::post;
29
30use crate::axum::error::Error;
31use crate::axum::extract::AgUiInput;
32use crate::axum::respond::SseResponse;
33
34/// Builds one transformer, once per run.
35type TransformerFactory = Arc<dyn Fn(&mut TransformerChain) + Send + Sync>;
36
37/// An agent plus the per-run settings it is served with.
38///
39/// [`route_agui`](RouterExt::route_agui) mounts an agent with the defaults;
40/// build one of these and use [`route_agui_with`](RouterExt::route_agui_with)
41/// when you want to change them.
42///
43/// ```
44/// use ag_ui::axum::{AgentEndpoint, RouterExt};
45/// use ag_ui::server::FilterToolCalls;
46/// use axum::Router;
47/// use std::time::Duration;
48/// # use ag_ui::RunOutcome;
49/// # use ag_ui::server::{Agent, Result, RunContext};
50/// # struct CartAgent;
51/// # impl Agent for CartAgent {
52/// #     type State = ();
53/// #     async fn run(&self, _ctx: &mut RunContext<()>) -> Result<RunOutcome> { Ok(RunOutcome::Success) }
54/// # }
55///
56/// let endpoint = AgentEndpoint::new(CartAgent)
57///     .transformer(|| FilterToolCalls::deny(["internal_debug"]))
58///     .keep_alive(Duration::from_secs(15));
59///
60/// let app: Router = Router::new().route_agui_with("/agent", endpoint);
61/// # let _ = app;
62/// ```
63pub struct AgentEndpoint<A> {
64    agent: Arc<A>,
65    transformers: Vec<TransformerFactory>,
66    echo_input: bool,
67    keep_alive: Option<Duration>,
68}
69
70impl<A> AgentEndpoint<A> {
71    /// Wraps an agent with the default settings.
72    pub fn new(agent: A) -> Self {
73        Self {
74            agent: Arc::new(agent),
75            transformers: Vec::new(),
76            echo_input: false,
77            keep_alive: None,
78        }
79    }
80
81    /// Appends a transformer to every run's chain.
82    ///
83    /// # Why a closure and not a transformer
84    ///
85    /// A [`StreamTransformer`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/server/transform/trait.StreamTransformer.html) takes `&mut self` because a useful one is a
86    /// state machine — [`FilterToolCalls`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/server/transform/struct.FilterToolCalls.html)
87    /// remembers which call ids it dropped. One instance shared across
88    /// concurrent runs would leak one run's state into another, so the endpoint
89    /// stores the recipe and builds a fresh chain per request.
90    #[must_use]
91    pub fn transformer<F, T>(mut self, factory: F) -> Self
92    where
93        F: Fn() -> T + Send + Sync + 'static,
94        T: StreamTransformer + 'static,
95    {
96        self.transformers
97            .push(Arc::new(move |chain: &mut TransformerChain| {
98                chain.push(factory());
99            }));
100        self
101    }
102
103    /// Echoes the request body back on `RUN_STARTED`.
104    ///
105    /// See [`Runner::echo_input`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/server/run/struct.Runner.html#method.echo_input). Off by default — it is the largest payload
106    /// in the protocol.
107    #[must_use]
108    pub fn echo_input(mut self, echo: bool) -> Self {
109        self.echo_input = echo;
110        self
111    }
112
113    /// Sends an SSE comment whenever a run produces nothing for `interval`.
114    ///
115    /// See [`SseResponse::keep_alive`]. Off by default.
116    #[must_use]
117    pub fn keep_alive(mut self, interval: Duration) -> Self {
118        self.keep_alive = Some(interval);
119        self
120    }
121
122    /// A fresh chain for one run.
123    fn chain(&self) -> TransformerChain {
124        let mut chain = TransformerChain::new();
125        for factory in &self.transformers {
126            factory(&mut chain);
127        }
128        chain
129    }
130}
131
132impl<A: Agent + 'static> AgentEndpoint<A> {
133    /// Serves one request: negotiate, decode, run, stream.
134    ///
135    /// Negotiation is reported first. A request that fails both checks has
136    /// nothing in common with this endpoint at all, and saying so is more use
137    /// to the caller than a note about a body that was never going to be
138    /// answered.
139    async fn serve(&self, headers: HeaderMap, input: Result<AgUiInput, Error>) -> Response {
140        let accept = headers
141            .get(header::ACCEPT)
142            .map(|value| String::from_utf8_lossy(value.as_bytes()));
143
144        let response = match SseResponse::negotiate(accept.as_deref()) {
145            Ok(response) => response,
146            Err(error) => return error.into_response(),
147        };
148        let input = match input {
149            Ok(AgUiInput(input)) => input,
150            Err(error) => return error.into_response(),
151        };
152
153        let runner = Runner::new(Arc::clone(&self.agent))
154            .transformers(self.chain())
155            .echo_input(self.echo_input);
156
157        // The token has to come off the runner before `run` consumes it.
158        let mut response = response.cancellation(runner.cancellation_token());
159        if let Some(interval) = self.keep_alive {
160            response = response.keep_alive(interval);
161        }
162        response.stream(runner.run(input))
163    }
164}
165
166impl<A> std::fmt::Debug for AgentEndpoint<A> {
167    /// Describes the settings, not the agent: an agent holds an LLM client and
168    /// a database handle, and neither is `Debug`.
169    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170        f.debug_struct("AgentEndpoint")
171            .field("agent", &std::any::type_name::<A>())
172            .field("transformers", &self.transformers.len())
173            .field("echo_input", &self.echo_input)
174            .field("keep_alive", &self.keep_alive)
175            .finish()
176    }
177}
178
179/// Mounts AG-UI agents on an [`axum::Router`].
180///
181/// # What it does to the router
182///
183/// `route_agui(path, agent)` is `route(path, post(handler))` and nothing else:
184/// the endpoint composes with the router's other routes, with `nest`, `merge`
185/// and `fallback`, and with any layer applied before or after it. A `GET` on
186/// the path still gets axum's own `405`.
187///
188/// # The state parameter
189///
190/// The handler reads only the request, so it is a `Handler<_, S>` for **every**
191/// router state `S`. Mounting an agent therefore places no constraint on `S`
192/// beyond axum's own `Clone + Send + Sync + 'static`, and works the same in a
193/// `Router<()>` and in a `Router<AppState>` — including before
194/// [`with_state`](axum::Router::with_state) is called.
195///
196/// An agent that needs values from the router state should capture them when it
197/// is constructed (`CartAgent::new(state.catalog.clone())`). Extracting `State`
198/// inside the AG-UI handler would tie this crate's one-line mount to a single
199/// application's state type, which is the opposite of what it is for.
200pub trait RouterExt<S>: Sized {
201    /// Mounts `agent` as a `POST` endpoint at `path`.
202    ///
203    /// The endpoint answers with `text/event-stream`, cancels the run when the
204    /// client disconnects, and refuses a request it cannot answer with a `4xx`.
205    #[must_use]
206    fn route_agui<A>(self, path: &str, agent: A) -> Self
207    where
208        A: Agent + 'static,
209    {
210        self.route_agui_with(path, AgentEndpoint::new(agent))
211    }
212
213    /// Mounts a configured [`AgentEndpoint`] as a `POST` endpoint at `path`.
214    #[must_use]
215    fn route_agui_with<A>(self, path: &str, endpoint: AgentEndpoint<A>) -> Self
216    where
217        A: Agent + 'static;
218}
219
220impl<S> RouterExt<S> for Router<S>
221where
222    S: Clone + Send + Sync + 'static,
223{
224    fn route_agui_with<A>(self, path: &str, endpoint: AgentEndpoint<A>) -> Self
225    where
226        A: Agent + 'static,
227    {
228        // One `Arc` for the endpoint, cloned per request; the agent itself is
229        // never cloned. `Arc<A>: Agent` is what lets the runner take it by
230        // value.
231        let endpoint = Arc::new(endpoint);
232        self.route(
233            path,
234            post(move |headers: HeaderMap, input: Result<AgUiInput, Error>| {
235                let endpoint = Arc::clone(&endpoint);
236                async move { endpoint.serve(headers, input).await }
237            }),
238        )
239    }
240}
ag_ui/axum/router.rs

ag_ui/axum/
router.rs
