ag_ui/axum/
extract.rs

1//! Reading a [`RunAgentInput`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/input/struct.RunAgentInput.html) out of the request.
2//!
3//! [`AgUiInput`] is a plain axum extractor, so an agent mounted by
4//! [`route_agui`](crate::axum::RouterExt::route_agui) and a hand-written handler that
5//! needs to look at the request first (auth, tenant routing, a
6//! [`Path`](axum::extract::Path) segment naming which agent to run) parse the
7//! body exactly the same way.
8//!
9//! ```
10//! use ag_ui::axum::AgUiInput;
11//! use axum::Router;
12//! use axum::extract::Path;
13//! use axum::routing::post;
14//!
15//! async fn handler(Path(agent): Path<String>, AgUiInput(input): AgUiInput) -> String {
16//!     format!("{agent} runs thread {}", input.thread_id)
17//! }
18//!
19//! let app: Router = Router::new().route("/agents/{agent}", post(handler));
20//! # let _ = app;
21//! ```
22//!
23//! Every failure is a `4xx` carrying a message that names the problem — see
24//! [`Error`]. Nothing here can panic on a hostile body: the size cap is axum's
25//! [`DefaultBodyLimit`](axum::extract::DefaultBodyLimit), and past that it is
26//! `serde_json` on a `&[u8]`.
27
28use crate::RunAgentInput;
29use axum::body::Bytes;
30use axum::extract::{FromRequest, Request};
31use axum::http::{HeaderMap, header};
32
33use crate::axum::error::{Error, Result};
34
35/// The AG-UI run request, extracted from a JSON body.
36///
37/// Rejects with [`Error`], which renders as a `4xx` and a JSON body naming what
38/// was wrong.
39#[derive(Clone, Debug, PartialEq)]
40pub struct AgUiInput(pub RunAgentInput);
41
42impl AgUiInput {
43    /// Unwraps the input.
44    pub fn into_inner(self) -> RunAgentInput {
45        self.0
46    }
47}
48
49impl<S> FromRequest<S> for AgUiInput
50where
51    S: Send + Sync,
52{
53    type Rejection = Error;
54
55    async fn from_request(request: Request, state: &S) -> Result<Self> {
56        check_content_type(request.headers())?;
57        // Through axum's own extractor rather than `to_bytes`, so a
58        // `DefaultBodyLimit` layer the user applied still applies here.
59        let bytes = Bytes::from_request(request, state)
60            .await
61            .map_err(|rejection| Error::Body {
62                status: rejection.status(),
63                message: rejection.body_text(),
64            })?;
65        Ok(Self(decode(&bytes)?))
66    }
67}
68
69/// Parses a `RunAgentInput` from raw JSON bytes.
70///
71/// The transport-free half of the extractor, for tests and for callers that
72/// already hold the body.
73///
74/// ```
75/// # use ag_ui::axum::extract::decode;
76/// let input = decode(br#"{"threadId":"t","runId":"r","messages":[],"tools":[],"context":[]}"#)?;
77/// assert_eq!(input.run_id.as_str(), "r");
78///
79/// assert!(decode(b"").is_err());
80/// assert!(decode(b"{").is_err());
81/// # Ok::<(), ag_ui::axum::Error>(())
82/// ```
83pub fn decode(body: &[u8]) -> Result<RunAgentInput> {
84    // serde's message for an empty body is "EOF while parsing a value at line 1
85    // column 0", which reads like a truncated payload rather than no payload.
86    if body.iter().all(u8::is_ascii_whitespace) {
87        return Err(Error::EmptyBody);
88    }
89    Ok(serde_json::from_slice(body)?)
90}
91
92/// Refuses a body whose `Content-Type` claims to be something other than JSON.
93///
94/// An absent header passes: plenty of clients omit it, and a missing header is
95/// never a browser form post — those always send one of three
96/// non-JSON types, which is what makes this check a CSRF defence as well as a
97/// content check.
98fn check_content_type(headers: &HeaderMap) -> Result<()> {
99    let Some(value) = headers.get(header::CONTENT_TYPE) else {
100        return Ok(());
101    };
102    let found = String::from_utf8_lossy(value.as_bytes());
103    if is_json(&found) {
104        return Ok(());
105    }
106    Err(Error::ContentType {
107        found: found.into_owned(),
108    })
109}
110
111/// Whether a `Content-Type` value names JSON, parameters and all.
112fn is_json(value: &str) -> bool {
113    let essence = value.split(';').next().unwrap_or(value).trim();
114    essence.eq_ignore_ascii_case("application/json")
115        || essence.eq_ignore_ascii_case("text/json")
116        // `application/vnd.acme+json` and friends.
117        || essence
118            .rsplit_once('+')
119            .is_some_and(|(_, suffix)| suffix.eq_ignore_ascii_case("json"))
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use axum::http::StatusCode;
126
127    fn headers(content_type: &str) -> HeaderMap {
128        let mut headers = HeaderMap::new();
129        headers.insert(
130            header::CONTENT_TYPE,
131            content_type.parse().expect("a valid header value"),
132        );
133        headers
134    }
135
136    #[test]
137    fn json_content_types_are_accepted() {
138        for value in [
139            "application/json",
140            "application/json; charset=utf-8",
141            "APPLICATION/JSON",
142            "text/json",
143            "application/vnd.acme.run+json",
144        ] {
145            assert!(
146                check_content_type(&headers(value)).is_ok(),
147                "should accept {value:?}"
148            );
149        }
150    }
151
152    #[test]
153    fn a_missing_content_type_is_accepted() {
154        assert!(check_content_type(&HeaderMap::new()).is_ok());
155    }
156
157    #[test]
158    fn a_form_post_is_refused_with_415() {
159        let error = check_content_type(&headers("application/x-www-form-urlencoded"))
160            .expect_err("should refuse a form body");
161        assert_eq!(error.status(), StatusCode::UNSUPPORTED_MEDIA_TYPE);
162        assert!(
163            error.to_string().contains("x-www-form-urlencoded"),
164            "{error}"
165        );
166    }
167
168    #[test]
169    fn a_missing_field_names_itself() {
170        let error = decode(br#"{"threadId":"t","runId":"r"}"#).expect_err("should not decode");
171        assert_eq!(error.status(), StatusCode::BAD_REQUEST);
172        assert!(error.to_string().contains("messages"), "{error}");
173    }
174
175    #[test]
176    fn whitespace_only_bodies_read_as_empty() {
177        assert!(matches!(decode(b"  \n\t"), Err(Error::EmptyBody)));
178    }
179}
ag_ui/axum/extract.rs

ag_ui/axum/
extract.rs
