ag_ui/server/agent.rs
1//! The boundary: implement [`Agent`], get an AG-UI endpoint.
2//!
3//! This crate depends on no LLM client. The .NET SDK can build on
4//! `Microsoft.Extensions.AI` because .NET has one blessed chat abstraction;
5//! Rust has `async-openai`, `rig-core` and `genai` with no winner, so binding
6//! to any of them would make this crate useless to two thirds of the ecosystem.
7//! A framework integration is an `impl Agent for …` in its own crate.
8
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::Arc;
12
13use crate::RunOutcome;
14use serde::Serialize;
15use serde::de::DeserializeOwned;
16
17use crate::server::context::RunContext;
18use crate::server::error::Result;
19
20/// What a run's shared state must be.
21///
22/// A blanket implementation covers every type that qualifies; you never write
23/// `impl AgentState`. Use `()` when the agent keeps no state.
24pub trait AgentState: Serialize + DeserializeOwned + Default + Send {}
25
26impl<T> AgentState for T where T: Serialize + DeserializeOwned + Default + Send {}
27
28/// An agent that can serve one AG-UI run.
29///
30/// This is the *hosting* side of the word. The consuming side —
31/// a handle onto somebody else's agent — is
32/// [`crate::client::RemoteAgent`], deliberately spelled differently so that an
33/// agent which calls another agent can import both.
34///
35/// [`crate::client::RemoteAgent`]: https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/client/agent/struct.RemoteAgent.html
36///
37/// ```
38/// use ag_ui::RunOutcome;
39/// use ag_ui::server::{Agent, Result, RunContext};
40///
41/// struct Echo;
42///
43/// impl Agent for Echo {
44/// type State = ();
45///
46/// async fn run(&self, ctx: &mut RunContext<()>) -> Result<RunOutcome> {
47/// let mut message = ctx.assistant_message()?;
48/// message.delta("you said something")?;
49/// message.end()?;
50/// Ok(RunOutcome::Success)
51/// }
52/// }
53/// ```
54///
55/// # Why `async fn` and not `#[async_trait]`
56///
57/// The trait uses a native `-> impl Future + Send` return (an RPITIT), so
58/// implementations are plain `async fn` with no macro, no `Box::pin` per call
59/// and no allocation. The cost is that the trait is not `dyn`-compatible; when
60/// you need `Box<dyn …>` — a registry of agents behind one endpoint, say —
61/// [`DynAgent`] is the boxed adapter, and [`BoxAgent`] implements `Agent`
62/// again, so the driver takes it like any other.
63///
64/// # Why `&mut RunContext` and not `RunContext`
65///
66/// The driver has to emit `RUN_FINISHED` or `RUN_ERROR` *after* the agent
67/// returns, through the same transformer chain and the same ordering verifier
68/// the agent used. Handing the context over by value would drop both with the
69/// agent's last statement.
70pub trait Agent: Send + Sync {
71 /// The run's shared state, deserialized from
72 /// [`RunAgentInput::state`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/input/struct.RunAgentInput.html#structfield.state) and published
73 /// through [`RunContext::set_state`].
74 type State: AgentState;
75
76 /// Serves one run.
77 ///
78 /// Returning `Ok(RunOutcome::Success)` finishes the run;
79 /// `Ok(RunOutcome::Interrupt { .. })` pauses it for human input; `Err`
80 /// becomes a `RUN_ERROR` event, never a panic and never a truncated
81 /// stream.
82 fn run(
83 &self,
84 ctx: &mut RunContext<Self::State>,
85 ) -> impl Future<Output = Result<RunOutcome>> + Send;
86}
87
88/// The `dyn`-compatible form of [`Agent`].
89///
90/// Implemented for every `Agent`, so `Box::new(my_agent) as BoxAgent<_>` just
91/// works. The only difference is one boxed future per run.
92pub trait DynAgent: Send + Sync {
93 /// The run's shared state — see [`Agent::State`].
94 type State: AgentState;
95
96 /// Serves one run, boxing the future so the trait stays object-safe.
97 fn run_boxed<'a>(
98 &'a self,
99 ctx: &'a mut RunContext<Self::State>,
100 ) -> Pin<Box<dyn Future<Output = Result<RunOutcome>> + Send + 'a>>;
101}
102
103impl<A: Agent> DynAgent for A {
104 type State = A::State;
105
106 fn run_boxed<'a>(
107 &'a self,
108 ctx: &'a mut RunContext<Self::State>,
109 ) -> Pin<Box<dyn Future<Output = Result<RunOutcome>> + Send + 'a>> {
110 Box::pin(self.run(ctx))
111 }
112}
113
114/// A type-erased agent over state `S`.
115///
116/// ```
117/// # use ag_ui::RunOutcome;
118/// # use ag_ui::server::{Agent, BoxAgent, Result, RunContext};
119/// struct Fixed(&'static str);
120///
121/// impl Agent for Fixed {
122/// type State = ();
123/// async fn run(&self, ctx: &mut RunContext<()>) -> Result<RunOutcome> {
124/// ctx.say(self.0)?;
125/// Ok(RunOutcome::Success)
126/// }
127/// }
128///
129/// let agents: Vec<BoxAgent<()>> = vec![Box::new(Fixed("a")), Box::new(Fixed("b"))];
130/// assert_eq!(agents.len(), 2);
131/// ```
132pub type BoxAgent<S> = Box<dyn DynAgent<State = S>>;
133
134impl<S: AgentState> Agent for BoxAgent<S> {
135 type State = S;
136
137 async fn run(&self, ctx: &mut RunContext<Self::State>) -> Result<RunOutcome> {
138 (**self).run_boxed(ctx).await
139 }
140}
141
142impl<A: Agent> Agent for &A {
143 type State = A::State;
144
145 async fn run(&self, ctx: &mut RunContext<Self::State>) -> Result<RunOutcome> {
146 (**self).run(ctx).await
147 }
148}
149
150impl<A: Agent> Agent for Arc<A> {
151 type State = A::State;
152
153 async fn run(&self, ctx: &mut RunContext<Self::State>) -> Result<RunOutcome> {
154 (**self).run(ctx).await
155 }
156}