ag_ui/encode/
mod.rs

1//! Wire encoding for the event stream.
2//!
3//! [`EventStreamFormatter`] is the abstraction a transport implements.
4#![cfg_attr(
5    feature = "sse",
6    doc = "[`sse::SseFormatter`] is the one every AG-UI implementation supports."
7)]
8#![cfg_attr(
9    not(feature = "sse"),
10    doc = "`sse::SseFormatter` is the one every AG-UI implementation supports, and the",
11    doc = "`sse` feature is off in this build."
12)]
13//!
14//! ```
15//! # #[cfg(feature = "sse")] {
16//! use ag_ui::{Event, EventStreamFormatter, SseFormatter};
17//!
18//! let formatter = SseFormatter::new();
19//! let bytes = formatter.encode(&Event::text_message_end("msg-1")).unwrap();
20//! assert_eq!(
21//!     String::from_utf8(bytes).unwrap(),
22//!     "data: {\"type\":\"TEXT_MESSAGE_END\",\"messageId\":\"msg-1\"}\n\n"
23//! );
24//! # }
25//! ```
26
27#[cfg(feature = "protobuf")]
28pub mod protobuf;
29#[cfg(feature = "sse")]
30pub mod sse;
31
32use crate::error::{Error, Result};
33use crate::event::Event;
34
35/// Media type of a Server-Sent Events stream.
36pub const SSE_MEDIA_TYPE: &str = "text/event-stream";
37
38/// Media type of the AG-UI binary stream, as defined by the upstream
39/// `@ag-ui/proto` package.
40pub const PROTOBUF_MEDIA_TYPE: &str = "application/vnd.ag-ui.event+proto";
41
42#[cfg(all(feature = "sse", feature = "protobuf"))]
43const SUPPORTED: &[&str] = &[SSE_MEDIA_TYPE, PROTOBUF_MEDIA_TYPE];
44#[cfg(all(feature = "sse", not(feature = "protobuf")))]
45const SUPPORTED: &[&str] = &[SSE_MEDIA_TYPE];
46#[cfg(all(not(feature = "sse"), feature = "protobuf"))]
47const SUPPORTED: &[&str] = &[PROTOBUF_MEDIA_TYPE];
48
49/// Turns events into the bytes a transport puts on the wire.
50///
51/// Object-safe, so a server can pick a formatter once per connection and store
52/// it as `Box<dyn EventStreamFormatter>`.
53pub trait EventStreamFormatter {
54    /// The `Content-Type` to answer with.
55    fn content_type(&self) -> &'static str;
56
57    /// Encodes one event, framing included.
58    fn encode(&self, event: &Event) -> Result<Vec<u8>>;
59}
60
61/// The media types this build can emit, most preferred first.
62pub const fn supported_media_types() -> &'static [&'static str] {
63    SUPPORTED
64}
65
66/// Picks the media type to answer an `Accept` header with.
67///
68/// A missing or empty header is treated as `*/*`, per RFC 9110. Candidates are
69/// scored by quality value; ties go to this crate's own preference order, which
70/// puts SSE first because it is the interoperable default and the only fully
71/// implemented transport here. That differs from the TypeScript encoder, which
72/// upgrades a bare `*/*` to protobuf.
73///
74/// Returns [`Error::UnsupportedMediaType`] when the header excludes everything
75/// this build can emit — the case that deserves a `406`.
76///
77/// ```
78/// # use ag_ui::encode::{media_type, SSE_MEDIA_TYPE};
79/// assert_eq!(media_type(None).unwrap(), SSE_MEDIA_TYPE);
80/// assert_eq!(media_type(Some("text/event-stream")).unwrap(), SSE_MEDIA_TYPE);
81/// assert!(media_type(Some("application/xml")).is_err());
82/// ```
83pub fn media_type(accept: Option<&str>) -> Result<&'static str> {
84    let header = accept
85        .map(str::trim)
86        .filter(|value| !value.is_empty())
87        .unwrap_or("*/*");
88
89    let specs: Vec<AcceptSpec<'_>> = header.split(',').filter_map(AcceptSpec::parse).collect();
90
91    let mut best: Option<(&'static str, f32)> = None;
92    for candidate in SUPPORTED {
93        let quality = quality_of(candidate, &specs);
94        if quality > 0.0 && best.is_none_or(|(_, best_quality)| quality > best_quality) {
95            best = Some((candidate, quality));
96        }
97    }
98
99    best.map(|(media_type, _)| media_type)
100        .ok_or_else(|| Error::UnsupportedMediaType(header.to_owned()))
101}
102
103/// One entry of an `Accept` header: a media range and its quality value.
104struct AcceptSpec<'a> {
105    kind: &'a str,
106    subtype: &'a str,
107    quality: f32,
108}
109
110impl<'a> AcceptSpec<'a> {
111    fn parse(entry: &'a str) -> Option<Self> {
112        let mut parts = entry.split(';');
113        let (kind, subtype) = split_media_type(parts.next()?.trim())?;
114
115        // Only `q` is read. The other parameters affect specificity in the full
116        // RFC 9110 algorithm, but no media type this crate emits is
117        // parameterized, so they can never change the outcome here.
118        let quality = parts
119            .find_map(|param| {
120                let (key, value) = param.split_once('=')?;
121                key.trim().eq_ignore_ascii_case("q").then_some(value)
122            })
123            .and_then(|value| value.trim().trim_matches('"').parse::<f32>().ok())
124            .unwrap_or(1.0);
125
126        Some(Self {
127            kind,
128            subtype,
129            quality,
130        })
131    }
132
133    /// How specifically this range names `kind`/`subtype`, or `None` when it
134    /// does not match at all. Higher is more specific.
135    fn specificity(&self, kind: &str, subtype: &str) -> Option<u8> {
136        let mut score = 0;
137
138        if self.kind.eq_ignore_ascii_case(kind) {
139            score |= 4;
140        } else if self.kind != "*" {
141            return None;
142        }
143
144        if self.subtype.eq_ignore_ascii_case(subtype) {
145            score |= 2;
146        } else if self.subtype != "*" {
147            return None;
148        }
149
150        Some(score)
151    }
152}
153
154/// The quality the header assigns to `candidate`, from its most specific
155/// matching range. Zero means unacceptable.
156fn quality_of(candidate: &str, specs: &[AcceptSpec<'_>]) -> f32 {
157    let Some((kind, subtype)) = split_media_type(candidate) else {
158        return 0.0;
159    };
160
161    let mut best: Option<(u8, f32)> = None;
162    for spec in specs {
163        let Some(specificity) = spec.specificity(kind, subtype) else {
164            continue;
165        };
166        if best.is_none_or(|(best_specificity, _)| specificity > best_specificity) {
167            best = Some((specificity, spec.quality));
168        }
169    }
170
171    best.map_or(0.0, |(_, quality)| quality)
172}
173
174/// Splits `type/subtype`, ignoring any parameters.
175fn split_media_type(value: &str) -> Option<(&str, &str)> {
176    let value = value.split(';').next()?.trim();
177    let (kind, subtype) = value.split_once('/')?;
178    (!kind.is_empty() && !subtype.is_empty()).then_some((kind, subtype))
179}
ag_ui/encode/mod.rs

ag_ui/encode/
mod.rs
