Skip to main content

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}